Upload 7 files

spectral/experiment_2_manifold_structures.ipynb

@@ -0,0 +1,2164 @@
+{
+ "nbformat": 4,
+ "nbformat_minor": 5,
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.10.0"
+  },
+  "accelerator": "GPU",
+  "colab": {
+   "provenance": [],
+   "gpuType": "T4"
+  }
+ },
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# GeoLIP Hypersphere Experiments \u2014 Category 2: Manifold Structures\n",
+    "\n",
+    "**Objective:** Explore structured manifolds beyond uniform hyperspheres.\n",
+    "\n",
+    "This notebook implements Hopf fibrations, Grassmannians, Flag manifolds, vMF soft assignment, and optimal lattice constructions for anchor initialization."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Shared Setup \u2014 Install & Configure"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "!pip install tensorboard scikit-learn matplotlib kymatio --quiet"
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "%load_ext tensorboard\n",
+    "%matplotlib inline"
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Shared Utilities Code"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "\n",
+    "# \u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\n",
+    "# SHARED UTILITIES \u2014 GeoLIP Hypersphere Experiments\n",
+    "# \u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\n",
+    "import os, sys, time, math, json, datetime, warnings\n",
+    "import numpy as np\n",
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "import torch.nn.functional as F\n",
+    "from torch.utils.data import DataLoader\n",
+    "from torch.utils.tensorboard import SummaryWriter\n",
+    "import torchvision\n",
+    "import torchvision.transforms as T\n",
+    "import matplotlib\n",
+    "matplotlib.use(\"Agg\")\n",
+    "import matplotlib.pyplot as plt\n",
+    "from collections import defaultdict\n",
+    "\n",
+    "warnings.filterwarnings(\"ignore\", category=UserWarning)\n",
+    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
+    "print(f\"[DEVICE] {device}\" + (f\" \u2014 {torch.cuda.get_device_name()}\" if device.type == \"cuda\" else \"\"))\n",
+    "\n",
+    "# \u2500\u2500 CIFAR-10 Loaders \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
+    "CIFAR10_MEAN = (0.4914, 0.4822, 0.4465)\n",
+    "CIFAR10_STD  = (0.2470, 0.2435, 0.2616)\n",
+    "CIFAR10_CLASSES = (\"plane\",\"car\",\"bird\",\"cat\",\"deer\",\"dog\",\"frog\",\"horse\",\"ship\",\"truck\")\n",
+    "\n",
+    "def get_cifar10_loaders(batch_size=256, num_workers=2, augment=True):\n",
+    "    \"\"\"Return train_loader, val_loader for CIFAR-10.\"\"\"\n",
+    "    tf_train = T.Compose([\n",
+    "        T.RandomCrop(32, padding=4),\n",
+    "        T.RandomHorizontalFlip(),\n",
+    "        T.ToTensor(),\n",
+    "        T.Normalize(CIFAR10_MEAN, CIFAR10_STD),\n",
+    "    ]) if augment else T.Compose([T.ToTensor(), T.Normalize(CIFAR10_MEAN, CIFAR10_STD)])\n",
+    "    tf_val = T.Compose([T.ToTensor(), T.Normalize(CIFAR10_MEAN, CIFAR10_STD)])\n",
+    "\n",
+    "    train_ds = torchvision.datasets.CIFAR10(root=\"./data\", train=True,  download=True, transform=tf_train)\n",
+    "    val_ds   = torchvision.datasets.CIFAR10(root=\"./data\", train=False, download=True, transform=tf_val)\n",
+    "    train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True,  num_workers=num_workers, pin_memory=True, drop_last=True)\n",
+    "    val_loader   = DataLoader(val_ds,   batch_size=batch_size, shuffle=False, num_workers=num_workers, pin_memory=True)\n",
+    "    print(f\"[DATA] CIFAR-10: {len(train_ds)} train, {len(val_ds)} val, batch_size={batch_size}\")\n",
+    "    return train_loader, val_loader\n",
+    "\n",
+    "def get_cifar10_raw_loaders(batch_size=256, num_workers=2):\n",
+    "    \"\"\"Return loaders with only ToTensor (no normalization) for custom front-ends.\"\"\"\n",
+    "    tf = T.Compose([T.ToTensor()])\n",
+    "    train_ds = torchvision.datasets.CIFAR10(root=\"./data\", train=True,  download=True, transform=tf)\n",
+    "    val_ds   = torchvision.datasets.CIFAR10(root=\"./data\", train=False, download=True, transform=tf)\n",
+    "    train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True,  num_workers=num_workers, pin_memory=True, drop_last=True)\n",
+    "    val_loader   = DataLoader(val_ds,   batch_size=batch_size, shuffle=False, num_workers=num_workers, pin_memory=True)\n",
+    "    print(f\"[DATA] CIFAR-10 (raw): {len(train_ds)} train, {len(val_ds)} val\")\n",
+    "    return train_loader, val_loader\n",
+    "\n",
+    "# \u2500\u2500 L2 Normalization to Hypersphere \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
+    "def l2_normalize(x, dim=-1, eps=1e-8):\n",
+    "    \"\"\"Project to unit hypersphere S^{d-1}.\"\"\"\n",
+    "    return F.normalize(x, p=2, dim=dim, eps=eps)\n",
+    "\n",
+    "# \u2500\u2500 Constellation Module \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
+    "class Constellation(nn.Module):\n",
+    "    \"\"\"\n",
+    "    K anchors on S^{d-1}. Triangulation: compute distance from input to each anchor.\n",
+    "    Output: K-dim vector of (1 - cos_sim) distances = triangulation profile.\n",
+    "    \"\"\"\n",
+    "    def __init__(self, embed_dim, num_anchors=64, init=\"uniform\"):\n",
+    "        super().__init__()\n",
+    "        self.embed_dim = embed_dim\n",
+    "        self.num_anchors = num_anchors\n",
+    "        anchors = self._init_anchors(embed_dim, num_anchors, init)\n",
+    "        self.anchors = nn.Parameter(anchors, requires_grad=False)  # anchors are NOT learned\n",
+    "\n",
+    "    def _init_anchors(self, d, K, method):\n",
+    "        if method == \"uniform\":\n",
+    "            # QR-based uniform hypersphere initialization\n",
+    "            A = torch.randn(K, d)\n",
+    "            Q, _ = torch.linalg.qr(A.T)\n",
+    "            anchors = Q.T[:K]\n",
+    "            # If K > d, fill remaining with random unit vectors\n",
+    "            if K > d:\n",
+    "                extra = F.normalize(torch.randn(K - d, d), dim=-1)\n",
+    "                anchors = torch.cat([anchors, extra], dim=0)\n",
+    "            return F.normalize(anchors, dim=-1)\n",
+    "        elif method == \"random\":\n",
+    "            return F.normalize(torch.randn(K, d), dim=-1)\n",
+    "        else:\n",
+    "            return F.normalize(torch.randn(K, d), dim=-1)\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        \"\"\"x: (B, d) on S^{d-1} \u2192 (B, K) triangulation distances.\"\"\"\n",
+    "        x_norm = l2_normalize(x)\n",
+    "        cos_sim = torch.mm(x_norm, self.anchors.T)  # (B, K)\n",
+    "        distances = 1.0 - cos_sim  # in [0, 2]\n",
+    "        return distances\n",
+    "\n",
+    "    def push_anchors(self, embeddings, labels, num_classes=10, momentum=0.1):\n",
+    "        \"\"\"Push anchors toward class centroids (call periodically).\"\"\"\n",
+    "        with torch.no_grad():\n",
+    "            for c in range(num_classes):\n",
+    "                mask = (labels == c)\n",
+    "                if mask.sum() < 2:\n",
+    "                    continue\n",
+    "                centroid = l2_normalize(embeddings[mask].mean(dim=0, keepdim=True))\n",
+    "                # Assign anchors to this class based on nearest centroid\n",
+    "                start = c * (self.num_anchors // num_classes)\n",
+    "                end = start + (self.num_anchors // num_classes)\n",
+    "                self.anchors.data[start:end] = l2_normalize(\n",
+    "                    (1 - momentum) * self.anchors.data[start:end] + momentum * centroid\n",
+    "                )\n",
+    "\n",
+    "    def anchor_activation(self, x):\n",
+    "        \"\"\"Return fraction of anchors that are 'active' (nearest to at least one sample).\"\"\"\n",
+    "        x_norm = l2_normalize(x)\n",
+    "        cos_sim = torch.mm(x_norm, self.anchors.T)\n",
+    "        nearest = cos_sim.argmax(dim=1)\n",
+    "        active = len(nearest.unique())\n",
+    "        return active, self.num_anchors\n",
+    "\n",
+    "# \u2500\u2500 Patchwork Block \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
+    "class SquaredReLU(nn.Module):\n",
+    "    def forward(self, x):\n",
+    "        return F.relu(x) ** 2\n",
+    "\n",
+    "class PatchworkBlock(nn.Module):\n",
+    "    \"\"\"SquaredReLU + LayerNorm block for post-constellation processing.\"\"\"\n",
+    "    def __init__(self, dim):\n",
+    "        super().__init__()\n",
+    "        self.net = nn.Sequential(\n",
+    "            nn.Linear(dim, dim),\n",
+    "            SquaredReLU(),\n",
+    "            nn.LayerNorm(dim),\n",
+    "            nn.Linear(dim, dim),\n",
+    "            SquaredReLU(),\n",
+    "            nn.LayerNorm(dim),\n",
+    "        )\n",
+    "    def forward(self, x):\n",
+    "        return x + self.net(x)  # residual\n",
+    "\n",
+    "class Patchwork(nn.Module):\n",
+    "    \"\"\"Stack of PatchworkBlocks.\"\"\"\n",
+    "    def __init__(self, dim, num_blocks=2):\n",
+    "        super().__init__()\n",
+    "        self.blocks = nn.Sequential(*[PatchworkBlock(dim) for _ in range(num_blocks)])\n",
+    "    def forward(self, x):\n",
+    "        return self.blocks(x)\n",
+    "\n",
+    "# \u2500\u2500 Classifier Head \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
+    "class ConstellationClassifier(nn.Module):\n",
+    "    \"\"\"Full pipeline: encoder output \u2192 constellation \u2192 patchwork \u2192 classify.\"\"\"\n",
+    "    def __init__(self, embed_dim, num_anchors=64, num_classes=10, patchwork_blocks=2):\n",
+    "        super().__init__()\n",
+    "        self.constellation = Constellation(embed_dim, num_anchors)\n",
+    "        self.patchwork = Patchwork(num_anchors, patchwork_blocks)\n",
+    "        self.head = nn.Linear(num_anchors, num_classes)\n",
+    "\n",
+    "    def forward(self, embeddings):\n",
+    "        tri = self.constellation(embeddings)       # (B, K)\n",
+    "        pw  = self.patchwork(tri)                   # (B, K)\n",
+    "        logits = self.head(pw)                      # (B, 10)\n",
+    "        return logits, tri\n",
+    "\n",
+    "    def count_params(self):\n",
+    "        return sum(p.numel() for p in self.parameters() if p.requires_grad)\n",
+    "\n",
+    "# \u2500\u2500 Losses \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
+    "class CayleyMengerVolumeLoss(nn.Module):\n",
+    "    \"\"\"CV loss: encourage constellation volume in 0.18-0.25 band.\"\"\"\n",
+    "    def __init__(self, target_low=0.18, target_high=0.25):\n",
+    "        super().__init__()\n",
+    "        self.target_low = target_low\n",
+    "        self.target_high = target_high\n",
+    "\n",
+    "    def forward(self, tri):\n",
+    "        \"\"\"tri: (B, K) triangulation distances.\"\"\"\n",
+    "        # Approximate volume via determinant of distance submatrix\n",
+    "        B, K = tri.shape\n",
+    "        # Use std of triangulation as proxy for volume spread\n",
+    "        vol = tri.std(dim=1).mean()\n",
+    "        # Penalty if outside target band\n",
+    "        below = F.relu(self.target_low - vol)\n",
+    "        above = F.relu(vol - self.target_high)\n",
+    "        return below + above, vol.item()\n",
+    "\n",
+    "class InfoNCELoss(nn.Module):\n",
+    "    \"\"\"Self-supervised contrastive loss on augmented views.\"\"\"\n",
+    "    def __init__(self, temperature=0.07):\n",
+    "        super().__init__()\n",
+    "        self.temperature = temperature\n",
+    "\n",
+    "    def forward(self, z1, z2):\n",
+    "        \"\"\"z1, z2: (B, d) L2-normalized embeddings of two augmented views.\"\"\"\n",
+    "        z1 = l2_normalize(z1)\n",
+    "        z2 = l2_normalize(z2)\n",
+    "        B = z1.shape[0]\n",
+    "        z = torch.cat([z1, z2], dim=0)  # (2B, d)\n",
+    "        sim = torch.mm(z, z.T) / self.temperature  # (2B, 2B)\n",
+    "        # Mask out self-similarity\n",
+    "        mask = ~torch.eye(2*B, dtype=bool, device=z.device)\n",
+    "        sim = sim.masked_fill(~mask, -1e9)\n",
+    "        # Positive pairs: (i, i+B) and (i+B, i)\n",
+    "        labels = torch.cat([torch.arange(B, 2*B), torch.arange(B)]).to(z.device)\n",
+    "        loss = F.cross_entropy(sim, labels)\n",
+    "        return loss\n",
+    "\n",
+    "class ClassConditionalInfoNCE(nn.Module):\n",
+    "    \"\"\"Class-conditional contrastive: positive = same class, not same image.\"\"\"\n",
+    "    def __init__(self, temperature=0.07):\n",
+    "        super().__init__()\n",
+    "        self.temperature = temperature\n",
+    "\n",
+    "    def forward(self, z, labels):\n",
+    "        z = l2_normalize(z)\n",
+    "        B = z.shape[0]\n",
+    "        sim = torch.mm(z, z.T) / self.temperature\n",
+    "        # Mask: positive if same class, negative otherwise\n",
+    "        label_mask = (labels.unsqueeze(0) == labels.unsqueeze(1)).float()\n",
+    "        label_mask.fill_diagonal_(0)  # exclude self\n",
+    "        # For each sample, compute log-softmax over all non-self, weighted by positives\n",
+    "        logits = sim - torch.diag(torch.full((B,), 1e9, device=z.device))\n",
+    "        pos_mask = label_mask\n",
+    "        neg_mask = 1 - label_mask\n",
+    "        neg_mask.fill_diagonal_(0)\n",
+    "        # InfoNCE: -log(sum_pos exp(sim) / sum_all exp(sim))\n",
+    "        exp_sim = torch.exp(logits)\n",
+    "        pos_sum = (exp_sim * pos_mask).sum(dim=1)\n",
+    "        all_sum = (exp_sim * (pos_mask + neg_mask)).sum(dim=1)\n",
+    "        loss = -torch.log(pos_sum / (all_sum + 1e-8) + 1e-8).mean()\n",
+    "        return loss\n",
+    "\n",
+    "# \u2500\u2500 TensorBoard Helpers \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
+    "def make_writer(experiment_id, base_dir=\"runs\"):\n",
+    "    \"\"\"Create a TensorBoard SummaryWriter with timestamped run directory.\"\"\"\n",
+    "    timestamp = datetime.datetime.now().strftime(\"%Y%m%d_%H%M%S\")\n",
+    "    log_dir = os.path.join(base_dir, f\"{experiment_id}_{timestamp}\")\n",
+    "    writer = SummaryWriter(log_dir=log_dir)\n",
+    "    print(f\"[TB] Logging to {log_dir}\")\n",
+    "    return writer\n",
+    "\n",
+    "def log_epoch_metrics(writer, epoch, metrics_dict, prefix=\"\"):\n",
+    "    \"\"\"Log a dict of metrics to TensorBoard.\"\"\"\n",
+    "    for key, val in metrics_dict.items():\n",
+    "        tag = f\"{prefix}/{key}\" if prefix else key\n",
+    "        writer.add_scalar(tag, val, epoch)\n",
+    "\n",
+    "def log_embedding_viz(writer, embeddings, labels, epoch, tag=\"embeddings\"):\n",
+    "    \"\"\"Log embeddings for TensorBoard projector.\"\"\"\n",
+    "    meta = [CIFAR10_CLASSES[l] for l in labels]\n",
+    "    writer.add_embedding(embeddings, metadata=meta, global_step=epoch, tag=tag)\n",
+    "\n",
+    "def log_anchor_histogram(writer, constellation, epoch):\n",
+    "    \"\"\"Log anchor norm distribution.\"\"\"\n",
+    "    norms = constellation.anchors.data.norm(dim=1)\n",
+    "    writer.add_histogram(\"anchors/norms\", norms, epoch)\n",
+    "    cos_sim_matrix = torch.mm(constellation.anchors.data, constellation.anchors.data.T)\n",
+    "    off_diag = cos_sim_matrix[~torch.eye(cos_sim_matrix.shape[0], dtype=bool)].flatten()\n",
+    "    writer.add_histogram(\"anchors/inter_cosine_sim\", off_diag, epoch)\n",
+    "\n",
+    "def log_confusion_matrix(writer, all_preds, all_labels, epoch):\n",
+    "    \"\"\"Log confusion matrix as image to TensorBoard.\"\"\"\n",
+    "    from sklearn.metrics import confusion_matrix\n",
+    "    cm = confusion_matrix(all_labels, all_preds, labels=range(10))\n",
+    "    fig, ax = plt.subplots(1, 1, figsize=(8, 8))\n",
+    "    im = ax.imshow(cm, interpolation=\"nearest\", cmap=plt.cm.Blues)\n",
+    "    ax.set_title(f\"Confusion Matrix \u2014 Epoch {epoch}\")\n",
+    "    ax.set_xlabel(\"Predicted\")\n",
+    "    ax.set_ylabel(\"True\")\n",
+    "    tick_marks = range(10)\n",
+    "    ax.set_xticks(tick_marks)\n",
+    "    ax.set_xticklabels(CIFAR10_CLASSES, rotation=45, ha=\"right\")\n",
+    "    ax.set_yticks(tick_marks)\n",
+    "    ax.set_yticklabels(CIFAR10_CLASSES)\n",
+    "    for i in range(10):\n",
+    "        for j in range(10):\n",
+    "            ax.text(j, i, str(cm[i, j]), ha=\"center\", va=\"center\",\n",
+    "                    color=\"white\" if cm[i, j] > cm.max() / 2 else \"black\", fontsize=7)\n",
+    "    fig.tight_layout()\n",
+    "    writer.add_figure(\"confusion_matrix\", fig, epoch)\n",
+    "    plt.close(fig)\n",
+    "\n",
+    "def log_triangulation_stats(writer, tri, epoch):\n",
+    "    \"\"\"Log triangulation distance statistics.\"\"\"\n",
+    "    writer.add_scalar(\"tri/mean\", tri.mean().item(), epoch)\n",
+    "    writer.add_scalar(\"tri/std\", tri.std().item(), epoch)\n",
+    "    writer.add_scalar(\"tri/min\", tri.min().item(), epoch)\n",
+    "    writer.add_scalar(\"tri/max\", tri.max().item(), epoch)\n",
+    "    writer.add_histogram(\"tri/distances\", tri.detach().cpu(), epoch)\n",
+    "\n",
+    "def log_scatter_2d(writer, embeddings, labels, epoch, tag=\"scatter_2d\", method=\"pca\"):\n",
+    "    \"\"\"Log 2D scatter plot of embeddings projected via PCA or t-SNE.\"\"\"\n",
+    "    emb_np = embeddings.detach().cpu().numpy()\n",
+    "    lab_np = labels.detach().cpu().numpy() if torch.is_tensor(labels) else np.array(labels)\n",
+    "    if method == \"pca\":\n",
+    "        from sklearn.decomposition import PCA\n",
+    "        proj = PCA(n_components=2).fit_transform(emb_np)\n",
+    "    else:\n",
+    "        from sklearn.manifold import TSNE\n",
+    "        proj = TSNE(n_components=2, perplexity=30, n_iter=500).fit_transform(emb_np)\n",
+    "    fig, ax = plt.subplots(1, 1, figsize=(8, 8))\n",
+    "    for c in range(10):\n",
+    "        mask = lab_np == c\n",
+    "        ax.scatter(proj[mask, 0], proj[mask, 1], s=3, alpha=0.5, label=CIFAR10_CLASSES[c])\n",
+    "    ax.legend(markerscale=5, fontsize=8)\n",
+    "    ax.set_title(f\"{tag} ({method.upper()}) \u2014 Epoch {epoch}\")\n",
+    "    fig.tight_layout()\n",
+    "    writer.add_figure(f\"{tag}/{method}\", fig, epoch)\n",
+    "    plt.close(fig)\n",
+    "\n",
+    "# \u2500\u2500 Training Utilities \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
+    "def train_one_epoch(model, loader, optimizer, criterion, device, epoch=0,\n",
+    "                    scheduler=None, extra_losses=None):\n",
+    "    \"\"\"\n",
+    "    Generic training loop for one epoch.\n",
+    "    model: callable(batch) \u2192 (logits, embeddings_on_sphere)\n",
+    "    extra_losses: list of (name, loss_fn, weight) where loss_fn(embeddings) \u2192 scalar\n",
+    "    \"\"\"\n",
+    "    model.train()\n",
+    "    total_loss = 0\n",
+    "    total_ce = 0\n",
+    "    correct = 0\n",
+    "    total = 0\n",
+    "    loss_components = defaultdict(float)\n",
+    "\n",
+    "    for batch_idx, (images, labels) in enumerate(loader):\n",
+    "        images, labels = images.to(device), labels.to(device)\n",
+    "        optimizer.zero_grad()\n",
+    "\n",
+    "        logits, embeddings = model(images)\n",
+    "        ce_loss = criterion(logits, labels)\n",
+    "        loss = ce_loss\n",
+    "        loss_components[\"CE\"] += ce_loss.item()\n",
+    "\n",
+    "        if extra_losses:\n",
+    "            for name, loss_fn, weight in extra_losses:\n",
+    "                if name == \"CV\":\n",
+    "                    el, vol_val = loss_fn(embeddings)\n",
+    "                    loss_components[\"CV_vol\"] += vol_val\n",
+    "                elif name == \"InfoNCE_class\":\n",
+    "                    el = loss_fn(embeddings, labels)\n",
+    "                else:\n",
+    "                    el = loss_fn(embeddings)\n",
+    "                loss = loss + weight * el\n",
+    "                loss_components[name] += el.item()\n",
+    "\n",
+    "        loss.backward()\n",
+    "        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)\n",
+    "        optimizer.step()\n",
+    "        if scheduler:\n",
+    "            scheduler.step()\n",
+    "\n",
+    "        total_loss += loss.item()\n",
+    "        _, predicted = logits.max(1)\n",
+    "        total += labels.size(0)\n",
+    "        correct += predicted.eq(labels).sum().item()\n",
+    "\n",
+    "    n_batches = len(loader)\n",
+    "    metrics = {\n",
+    "        \"loss/total\": total_loss / n_batches,\n",
+    "        \"loss/CE\": loss_components[\"CE\"] / n_batches,\n",
+    "        \"accuracy/train\": 100.0 * correct / total,\n",
+    "    }\n",
+    "    for k, v in loss_components.items():\n",
+    "        if k not in (\"CE\",):\n",
+    "            metrics[f\"loss/{k}\"] = v / n_batches\n",
+    "    return metrics\n",
+    "\n",
+    "@torch.no_grad()\n",
+    "def evaluate(model, loader, criterion, device, collect_embeddings=False):\n",
+    "    \"\"\"Evaluate model on validation set. Optionally collect embeddings.\"\"\"\n",
+    "    model.eval()\n",
+    "    total_loss = 0\n",
+    "    correct = 0\n",
+    "    total = 0\n",
+    "    all_preds = []\n",
+    "    all_labels = []\n",
+    "    all_embeddings = []\n",
+    "\n",
+    "    for images, labels in loader:\n",
+    "        images, labels = images.to(device), labels.to(device)\n",
+    "        logits, embeddings = model(images)\n",
+    "        loss = criterion(logits, labels)\n",
+    "        total_loss += loss.item()\n",
+    "        _, predicted = logits.max(1)\n",
+    "        total += labels.size(0)\n",
+    "        correct += predicted.eq(labels).sum().item()\n",
+    "        all_preds.extend(predicted.cpu().numpy())\n",
+    "        all_labels.extend(labels.cpu().numpy())\n",
+    "        if collect_embeddings:\n",
+    "            all_embeddings.append(embeddings.cpu())\n",
+    "\n",
+    "    n_batches = len(loader)\n",
+    "    metrics = {\n",
+    "        \"loss/val\": total_loss / n_batches,\n",
+    "        \"accuracy/val\": 100.0 * correct / total,\n",
+    "    }\n",
+    "    result = {\"metrics\": metrics, \"preds\": np.array(all_preds), \"labels\": np.array(all_labels)}\n",
+    "    if collect_embeddings:\n",
+    "        result[\"embeddings\"] = torch.cat(all_embeddings, dim=0)\n",
+    "    return result\n",
+    "\n",
+    "def run_experiment(experiment_id, model, train_loader, val_loader, device,\n",
+    "                   epochs=100, lr=1e-3, weight_decay=1e-4, extra_losses=None,\n",
+    "                   push_anchors_every=50, log_embeddings_every=25,\n",
+    "                   patience=20):\n",
+    "    \"\"\"\n",
+    "    Full experiment runner with TensorBoard logging, early stopping, and diagnostics.\n",
+    "    model must return (logits, sphere_embeddings).\n",
+    "    \"\"\"\n",
+    "    writer = make_writer(experiment_id)\n",
+    "    criterion = nn.CrossEntropyLoss()\n",
+    "    optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=weight_decay)\n",
+    "    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs)\n",
+    "\n",
+    "    best_val_acc = 0.0\n",
+    "    best_epoch = 0\n",
+    "    no_improve = 0\n",
+    "\n",
+    "    # Log model info\n",
+    "    n_params = sum(p.numel() for p in model.parameters() if p.requires_grad)\n",
+    "    writer.add_text(\"config/params\", f\"Learned params: {n_params:,}\")\n",
+    "    writer.add_text(\"config/experiment\", experiment_id)\n",
+    "    print(f\"\\n{'='*70}\")\n",
+    "    print(f\"[EXP] {experiment_id}\")\n",
+    "    print(f\"[EXP] Learned params: {n_params:,}, Epochs: {epochs}, LR: {lr}\")\n",
+    "    print(f\"{'='*70}\")\n",
+    "\n",
+    "    for epoch in range(1, epochs + 1):\n",
+    "        t0 = time.time()\n",
+    "\n",
+    "        # Train\n",
+    "        train_metrics = train_one_epoch(\n",
+    "            model, train_loader, optimizer, criterion, device, epoch,\n",
+    "            extra_losses=extra_losses\n",
+    "        )\n",
+    "\n",
+    "        # Evaluate\n",
+    "        collect_emb = (epoch % log_embeddings_every == 0) or (epoch == epochs)\n",
+    "        val_result = evaluate(model, val_loader, criterion, device, collect_embeddings=collect_emb)\n",
+    "        val_metrics = val_result[\"metrics\"]\n",
+    "\n",
+    "        # Combine metrics\n",
+    "        all_metrics = {**train_metrics, **val_metrics}\n",
+    "        all_metrics[\"lr\"] = optimizer.param_groups[0][\"lr\"]\n",
+    "        all_metrics[\"epoch_time\"] = time.time() - t0\n",
+    "        all_metrics[\"train_val_gap\"] = train_metrics[\"accuracy/train\"] - val_metrics[\"accuracy/val\"]\n",
+    "\n",
+    "        # Constellation diagnostics\n",
+    "        if hasattr(model, \"constellation\"):\n",
+    "            constellation = model.constellation\n",
+    "        elif hasattr(model, \"classifier\") and hasattr(model.classifier, \"constellation\"):\n",
+    "            constellation = model.classifier.constellation\n",
+    "        else:\n",
+    "            constellation = None\n",
+    "\n",
+    "        if constellation is not None:\n",
+    "            log_anchor_histogram(writer, constellation, epoch)\n",
+    "\n",
+    "            # Push anchors periodically\n",
+    "            if epoch % push_anchors_every == 0 and \"embeddings\" in val_result:\n",
+    "                emb = val_result[\"embeddings\"][:2000].to(device)\n",
+    "                lab = torch.tensor(val_result[\"labels\"][:2000]).to(device)\n",
+    "                constellation.push_anchors(emb, lab)\n",
+    "                active, total_anch = constellation.anchor_activation(emb)\n",
+    "                all_metrics[\"anchors/active\"] = active\n",
+    "                all_metrics[\"anchors/total\"] = total_anch\n",
+    "                all_metrics[\"anchors/utilization\"] = active / total_anch\n",
+    "                print(f\"  [ANCHOR PUSH] Active: {active}/{total_anch} ({100*active/total_anch:.1f}%)\")\n",
+    "\n",
+    "        # Log all metrics\n",
+    "        log_epoch_metrics(writer, epoch, all_metrics)\n",
+    "        scheduler.step()\n",
+    "\n",
+    "        # Log embeddings visualization\n",
+    "        if collect_emb and \"embeddings\" in val_result:\n",
+    "            emb = val_result[\"embeddings\"][:2000]\n",
+    "            lab = val_result[\"labels\"][:2000]\n",
+    "            log_scatter_2d(writer, emb, lab, epoch, tag=\"val_embeddings\", method=\"pca\")\n",
+    "            if epoch == epochs:\n",
+    "                log_confusion_matrix(writer, val_result[\"preds\"], val_result[\"labels\"], epoch)\n",
+    "                log_scatter_2d(writer, emb, lab, epoch, tag=\"val_embeddings_tsne\", method=\"tsne\")\n",
+    "\n",
+    "        # Log triangulation stats from a sample batch\n",
+    "        if constellation is not None and collect_emb and \"embeddings\" in val_result:\n",
+    "            sample_emb = val_result[\"embeddings\"][:500].to(device)\n",
+    "            tri_sample = constellation(sample_emb)\n",
+    "            log_triangulation_stats(writer, tri_sample, epoch)\n",
+    "\n",
+    "        # Best model tracking\n",
+    "        val_acc = val_metrics[\"accuracy/val\"]\n",
+    "        if val_acc > best_val_acc:\n",
+    "            best_val_acc = val_acc\n",
+    "            best_epoch = epoch\n",
+    "            no_improve = 0\n",
+    "            writer.add_scalar(\"best/val_acc\", best_val_acc, epoch)\n",
+    "            writer.add_scalar(\"best/epoch\", best_epoch, epoch)\n",
+    "        else:\n",
+    "            no_improve += 1\n",
+    "\n",
+    "        # Print progress\n",
+    "        if epoch % 5 == 0 or epoch == 1 or epoch == epochs:\n",
+    "            print(f\"  E{epoch:3d} | TrainAcc {train_metrics['accuracy/train']:5.1f}% | \"\n",
+    "                  f\"ValAcc {val_acc:5.1f}% | Loss {all_metrics['loss/total']:.4f} | \"\n",
+    "                  f\"Best {best_val_acc:.1f}% @E{best_epoch} | \"\n",
+    "                  f\"Gap {all_metrics['train_val_gap']:+.1f}% | \"\n",
+    "                  f\"{all_metrics['epoch_time']:.1f}s\")\n",
+    "\n",
+    "        # Early stopping\n",
+    "        if no_improve >= patience:\n",
+    "            print(f\"  [EARLY STOP] No improvement for {patience} epochs. Best: {best_val_acc:.1f}% @E{best_epoch}\")\n",
+    "            break\n",
+    "\n",
+    "    writer.close()\n",
+    "    print(f\"\\n[RESULT] {experiment_id}: Best Val Acc = {best_val_acc:.2f}% at epoch {best_epoch}\")\n",
+    "    print(f\"[RESULT] Learned params: {n_params:,}\")\n",
+    "    return {\"experiment\": experiment_id, \"best_val_acc\": best_val_acc, \"best_epoch\": best_epoch,\n",
+    "            \"params\": n_params, \"final_train_acc\": train_metrics[\"accuracy/train\"]}\n",
+    "\n",
+    "# \u2500\u2500 Johnson-Lindenstrauss Projection \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
+    "class JLProjection(nn.Module):\n",
+    "    \"\"\"Fixed random Gaussian projection (not learned).\"\"\"\n",
+    "    def __init__(self, input_dim, output_dim):\n",
+    "        super().__init__()\n",
+    "        self.register_buffer(\"W\", torch.randn(input_dim, output_dim) / math.sqrt(output_dim))\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        return x @ self.W\n",
+    "\n",
+    "# \u2500\u2500 Results Scoreboard \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
+    "class Scoreboard:\n",
+    "    \"\"\"Track results across experiments in a session.\"\"\"\n",
+    "    def __init__(self):\n",
+    "        self.results = []\n",
+    "\n",
+    "    def add(self, result_dict):\n",
+    "        self.results.append(result_dict)\n",
+    "        self._print_table()\n",
+    "\n",
+    "    def _print_table(self):\n",
+    "        print(f\"\\n{'='*80}\")\n",
+    "        print(f\"{'SCOREBOARD':^80}\")\n",
+    "        print(f\"{'='*80}\")\n",
+    "        print(f\"{'Experiment':<40} {'Val Acc':>8} {'Params':>10} {'Best E':>7}\")\n",
+    "        print(f\"{'-'*40} {'-'*8} {'-'*10} {'-'*7}\")\n",
+    "        for r in sorted(self.results, key=lambda x: x[\"best_val_acc\"], reverse=True):\n",
+    "            print(f\"{r['experiment']:<40} {r['best_val_acc']:>7.2f}% {r['params']:>10,} {r['best_epoch']:>7}\")\n",
+    "        print(f\"{'='*80}\\n\")\n",
+    "\n",
+    "    def save(self, path=\"scoreboard.json\"):\n",
+    "        with open(path, \"w\") as f:\n",
+    "            json.dump(self.results, f, indent=2)\n",
+    "        print(f\"[SCOREBOARD] Saved to {path}\")\n",
+    "\n",
+    "scoreboard = Scoreboard()\n",
+    "print(\"[SHARED] All utilities loaded successfully.\")\n",
+    ""
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Experiment 2.1a: Hopf Fibration S\u00b3 \u2192 S\u00b2"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "\n",
+    "Formula: h(z\u2081,z\u2082) = (2z\u0304\u2081z\u2082, |z\u2081|\u00b2\u2212|z\u2082|\u00b2) : S\u00b3 \u2192 S\u00b2\n",
+    "\n",
+    "Strategy:\n",
+    "- Encode 4-d feature vectors on S\u00b3\n",
+    "- Hopf project to S\u00b2 (coarse triangulation) + fiber coordinate (fine discrimination)\n",
+    "- Front-end: Scattering2D + JL projection to 4-d\n",
+    "- TensorBoard logging, 100 epochs, early stopping patience=20\n",
+    ""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "\n",
+    "from kymatio.torch import Scattering2D\n",
+    "import torch.optim.lr_scheduler as lr_scheduler\n",
+    "\n",
+    "# \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
+    "# FRONT-END: Scattering2D + JL to 4-d\n",
+    "# \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
+    "class Scattering2DFeatures(nn.Module):\n",
+    "    '''Scattering2D feature extractor for CIFAR-10.'''\n",
+    "    def __init__(self, out_dim=128):\n",
+    "        super().__init__()\n",
+    "        # Scattering2D for 32x32 images\n",
+    "        self.scattering = Scattering2D(J=2, shape=(32, 32))\n",
+    "        scat_out = self.compute_scat_size()\n",
+    "        self.proj = nn.Linear(scat_out, out_dim)\n",
+    "        print(f\"[SCAT] Scattering output size: {scat_out}, projected to {out_dim}\")\n",
+    "\n",
+    "    def compute_scat_size(self):\n",
+    "        '''Compute output size of Scattering2D.'''\n",
+    "        x_dummy = torch.randn(1, 3, 32, 32)\n",
+    "        with torch.no_grad():\n",
+    "            out = self.scattering(x_dummy)\n",
+    "        return out.shape[1] * out.shape[2] * out.shape[3]\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        # x: (B, 3, 32, 32)\n",
+    "        scats = []\n",
+    "        for c in range(3):\n",
+    "            s = self.scattering(x[:, c:c+1, :, :])  # (B, F, H, W)\n",
+    "            s = s.reshape(s.shape[0], -1)  # (B, F*H*W)\n",
+    "            scats.append(s)\n",
+    "        scat_feat = torch.cat(scats, dim=1)  # (B, 3*F*H*W)\n",
+    "        out = self.proj(scat_feat)  # (B, out_dim)\n",
+    "        return out\n",
+    "\n",
+    "class HopfFibrationS3S2Model(nn.Module):\n",
+    "    '''S\u00b3 \u2192 S\u00b2 via Hopf fibration (z\u2081, z\u2082 complex).'''\n",
+    "    def __init__(self, num_anchors=64):\n",
+    "        super().__init__()\n",
+    "        self.frontend = Scattering2DFeatures(out_dim=128)\n",
+    "        self.to_4d = nn.Linear(128, 4)  # Project to 4-d for S\u00b3\n",
+    "        self.constellation_s2 = Constellation(embed_dim=3, num_anchors=num_anchors, init=\"uniform\")\n",
+    "        self.fiber_encoder = nn.Linear(1, 16)  # Encode fiber coordinate\n",
+    "        self.patchwork = Patchwork(num_anchors + 16, num_blocks=2)\n",
+    "        self.head = nn.Linear(num_anchors + 16, 10)\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        # Frontend: scattering + projection\n",
+    "        feat = self.frontend(x)  # (B, 128)\n",
+    "        z_4d = self.to_4d(feat)  # (B, 4)\n",
+    "        z_norm = l2_normalize(z_4d)  # Normalize to S\u00b3\n",
+    "\n",
+    "        # Hopf projection: S\u00b3 \u2192 S\u00b2 (interpret as z\u2081, z\u2082 \u2208 \u2102)\n",
+    "        # z = [z\u2081.re, z\u2081.im, z\u2082.re, z\u2082.im]\n",
+    "        z1 = torch.complex(z_norm[:, 0], z_norm[:, 1])\n",
+    "        z2 = torch.complex(z_norm[:, 2], z_norm[:, 3])\n",
+    "\n",
+    "        # Hopf map: (2z\u0304\u2081z\u2082, |z\u2081|\u00b2\u2212|z\u2082|\u00b2)\n",
+    "        two_conj_z1_z2 = 2 * (z1.conj() * z2)\n",
+    "        z1_norm_sq = (z1.abs() ** 2).unsqueeze(-1)\n",
+    "        z2_norm_sq = (z2.abs() ** 2).unsqueeze(-1)\n",
+    "\n",
+    "        base_x = torch.cat([\n",
+    "            two_conj_z1_z2.real.unsqueeze(-1),\n",
+    "            two_conj_z1_z2.imag.unsqueeze(-1),\n",
+    "            (z1_norm_sq - z2_norm_sq)\n",
+    "        ], dim=1)  # (B, 3) on S\u00b2\n",
+    "        base_x = l2_normalize(base_x)\n",
+    "\n",
+    "        # Fiber coordinate: one of the complex arguments\n",
+    "        fiber = (z1.abs()).unsqueeze(-1)  # (B, 1)\n",
+    "\n",
+    "        # Triangulation on S\u00b2 + fiber encoding\n",
+    "        tri = self.constellation_s2(base_x)  # (B, num_anchors)\n",
+    "        fiber_enc = self.fiber_encoder(fiber)  # (B, 16)\n",
+    "\n",
+    "        # Combine\n",
+    "        combined = torch.cat([tri, fiber_enc], dim=1)  # (B, num_anchors+16)\n",
+    "        pw = self.patchwork(combined)\n",
+    "        logits = self.head(pw)\n",
+    "\n",
+    "        # Return logits and sphere embedding (for diagnostics)\n",
+    "        return logits, z_norm\n",
+    "\n",
+    "# \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
+    "# RUN EXPERIMENT 2.1a\n",
+    "# \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
+    "print(\"\\n[EXP 2.1a] Loading CIFAR-10...\")\n",
+    "train_loader, val_loader = get_cifar10_raw_loaders(batch_size=128, num_workers=2)\n",
+    "\n",
+    "model_2_1a = HopfFibrationS3S2Model(num_anchors=64).to(device)\n",
+    "print(f\"[EXP 2.1a] Model: {model_2_1a.count_params() if hasattr(model_2_1a, 'count_params') else 'N/A'} params\")\n",
+    "\n",
+    "result_2_1a = run_experiment(\n",
+    "    \"2.1a_HopfS3S2\",\n",
+    "    model_2_1a,\n",
+    "    train_loader,\n",
+    "    val_loader,\n",
+    "    device,\n",
+    "    epochs=100,\n",
+    "    lr=1e-3,\n",
+    "    patience=20\n",
+    ")\n",
+    "scoreboard.add(result_2_1a)\n",
+    ""
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Experiment 2.1b: Quaternionic Hopf S\u2077 \u2192 S\u2074"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "\n",
+    "Natural decomposition for 8-d embeddings (8-channel spectral).\n",
+    "Front-end produces 8-d embeddings \u2192 Hopf project to S\u2074 + fiber.\n",
+    ""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "\n",
+    "class HopfFibrationS7S4Model(nn.Module):\n",
+    "    '''S\u2077 \u2192 S\u2074 via quaternionic Hopf (q\u2081, q\u2082 quaternions).'''\n",
+    "    def __init__(self, num_anchors=64):\n",
+    "        super().__init__()\n",
+    "        self.frontend = Scattering2DFeatures(out_dim=256)\n",
+    "        self.to_8d = nn.Linear(256, 8)\n",
+    "        self.constellation_s4 = Constellation(embed_dim=5, num_anchors=num_anchors, init=\"uniform\")\n",
+    "        self.fiber_encoder = nn.Linear(3, 16)\n",
+    "        self.patchwork = Patchwork(num_anchors + 16, num_blocks=2)\n",
+    "        self.head = nn.Linear(num_anchors + 16, 10)\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        feat = self.frontend(x)  # (B, 256)\n",
+    "        z_8d = self.to_8d(feat)  # (B, 8)\n",
+    "        z_norm = l2_normalize(z_8d)  # Normalize to S\u2077\n",
+    "\n",
+    "        # Quaternionic Hopf: q\u2081 = (z\u2081, z\u2082), q\u2082 = (z\u2083, z\u2084) as complex pairs\n",
+    "        # Map to S\u2074: base = [2Re(conj(q\u2081)q\u2082), |q\u2081|\u00b2 - |q\u2082|\u00b2]\n",
+    "        q1_0 = torch.complex(z_norm[:, 0], z_norm[:, 1])\n",
+    "        q1_1 = torch.complex(z_norm[:, 2], z_norm[:, 3])\n",
+    "        q2_0 = torch.complex(z_norm[:, 4], z_norm[:, 5])\n",
+    "        q2_1 = torch.complex(z_norm[:, 6], z_norm[:, 7])\n",
+    "\n",
+    "        # Simplified quaternionic norm\n",
+    "        q1_norm_sq = (q1_0.abs() ** 2 + q1_1.abs() ** 2).unsqueeze(-1)\n",
+    "        q2_norm_sq = (q2_0.abs() ** 2 + q2_1.abs() ** 2).unsqueeze(-1)\n",
+    "\n",
+    "        # Quaternion multiplication (simplified)\n",
+    "        conj_q1_q2_0 = (q1_0.conj() * q2_0 - q1_1.conj() * q2_1)\n",
+    "        conj_q1_q2_1 = (q1_0.conj() * q2_1 + q1_1.conj() * q2_0)\n",
+    "\n",
+    "        base_x = torch.cat([\n",
+    "            conj_q1_q2_0.real.unsqueeze(-1),\n",
+    "            conj_q1_q2_0.imag.unsqueeze(-1),\n",
+    "            conj_q1_q2_1.real.unsqueeze(-1),\n",
+    "            conj_q1_q2_1.imag.unsqueeze(-1),\n",
+    "            (q1_norm_sq - q2_norm_sq)\n",
+    "        ], dim=1)  # (B, 5) on S\u2074\n",
+    "        base_x = l2_normalize(base_x)\n",
+    "\n",
+    "        # Fiber: three components\n",
+    "        fiber = torch.stack([\n",
+    "            (q1_0.abs()),\n",
+    "            (q1_1.abs()),\n",
+    "            (q2_0.abs())\n",
+    "        ], dim=1)  # (B, 3)\n",
+    "\n",
+    "        tri = self.constellation_s4(base_x)  # (B, num_anchors)\n",
+    "        fiber_enc = self.fiber_encoder(fiber)  # (B, 16)\n",
+    "        combined = torch.cat([tri, fiber_enc], dim=1)\n",
+    "        pw = self.patchwork(combined)\n",
+    "        logits = self.head(pw)\n",
+    "        return logits, z_norm\n",
+    "\n",
+    "# \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
+    "# RUN EXPERIMENT 2.1b\n",
+    "# \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
+    "print(\"\\n[EXP 2.1b] Setting up S\u2077 \u2192 S\u2074...\")\n",
+    "model_2_1b = HopfFibrationS7S4Model(num_anchors=64).to(device)\n",
+    "result_2_1b = run_experiment(\n",
+    "    \"2.1b_HopfS7S4\",\n",
+    "    model_2_1b,\n",
+    "    train_loader,\n",
+    "    val_loader,\n",
+    "    device,\n",
+    "    epochs=100,\n",
+    "    lr=1e-3,\n",
+    "    patience=20\n",
+    ")\n",
+    "scoreboard.add(result_2_1b)\n",
+    ""
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Experiment 2.1c: Hopf Foliation Anchor Init"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "\n",
+    "Replace uniform_hypersphere_init with Hopf-structured codes.\n",
+    "Compare anchor utilization vs random initialization.\n",
+    ""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "\n",
+    "class HopfAnchorConstellation(nn.Module):\n",
+    "    '''Constellation with Hopf-structured anchor initialization.'''\n",
+    "    def __init__(self, embed_dim, num_anchors=64):\n",
+    "        super().__init__()\n",
+    "        self.embed_dim = embed_dim\n",
+    "        self.num_anchors = num_anchors\n",
+    "        anchors = self._init_hopf_anchors(embed_dim, num_anchors)\n",
+    "        self.anchors = nn.Parameter(anchors, requires_grad=False)\n",
+    "\n",
+    "    def _init_hopf_anchors(self, d, K):\n",
+    "        '''Initialize anchors on S^{d-1} using Hopf decomposition.'''\n",
+    "        # For d=4 (S\u00b3), use Hopf fibration structure\n",
+    "        if d == 4:\n",
+    "            n_base = int(np.sqrt(K / 2))\n",
+    "            base_angles = np.linspace(0, 2*np.pi, n_base, endpoint=False)\n",
+    "            fiber_angles = np.linspace(0, np.pi, n_base, endpoint=False)\n",
+    "\n",
+    "            anchors = []\n",
+    "            for theta in base_angles:\n",
+    "                for phi in fiber_angles:\n",
+    "                    # S\u00b2 base coordinates\n",
+    "                    x = np.sin(theta) * np.cos(phi)\n",
+    "                    y = np.sin(theta) * np.sin(phi)\n",
+    "                    z = np.cos(theta)\n",
+    "                    # Lift to S\u00b3: use phi as fiber parameter\n",
+    "                    z1_real, z1_imag = np.cos(phi/2), np.sin(phi/2)\n",
+    "                    z2_real, z2_imag = np.sin(phi/2) * np.cos(theta), np.sin(phi/2) * np.sin(theta)\n",
+    "                    anchors.append([z1_real, z1_imag, z2_real, z2_imag])\n",
+    "            anchors = torch.tensor(anchors[:K], dtype=torch.float32)\n",
+    "        else:\n",
+    "            # Fallback: uniform\n",
+    "            A = torch.randn(K, d)\n",
+    "            Q, _ = torch.linalg.qr(A.T)\n",
+    "            anchors = Q.T[:K]\n",
+    "\n",
+    "        return F.normalize(anchors, dim=-1)\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        x_norm = l2_normalize(x)\n",
+    "        cos_sim = torch.mm(x_norm, self.anchors.T)\n",
+    "        distances = 1.0 - cos_sim\n",
+    "        return distances\n",
+    "\n",
+    "    def anchor_activation(self, x):\n",
+    "        x_norm = l2_normalize(x)\n",
+    "        cos_sim = torch.mm(x_norm, self.anchors.T)\n",
+    "        nearest = cos_sim.argmax(dim=1)\n",
+    "        active = len(nearest.unique())\n",
+    "        return active, self.num_anchors\n",
+    "\n",
+    "class HopfAnchorModel(nn.Module):\n",
+    "    '''S\u00b3 embedding with Hopf-structured anchors.'''\n",
+    "    def __init__(self, num_anchors=64):\n",
+    "        super().__init__()\n",
+    "        self.frontend = Scattering2DFeatures(out_dim=128)\n",
+    "        self.to_4d = nn.Linear(128, 4)\n",
+    "        self.constellation = HopfAnchorConstellation(embed_dim=4, num_anchors=num_anchors)\n",
+    "        self.patchwork = Patchwork(num_anchors, num_blocks=2)\n",
+    "        self.head = nn.Linear(num_anchors, 10)\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        feat = self.frontend(x)\n",
+    "        z_4d = self.to_4d(feat)\n",
+    "        z_norm = l2_normalize(z_4d)\n",
+    "        tri = self.constellation(z_norm)\n",
+    "        pw = self.patchwork(tri)\n",
+    "        logits = self.head(pw)\n",
+    "        return logits, z_norm\n",
+    "\n",
+    "# \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
+    "# COMPARE: Hopf-structured vs Random\n",
+    "# \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
+    "print(\"\\n[EXP 2.1c] Hopf-structured Anchors...\")\n",
+    "model_2_1c = HopfAnchorModel(num_anchors=64).to(device)\n",
+    "result_2_1c = run_experiment(\n",
+    "    \"2.1c_HopfAnchors\",\n",
+    "    model_2_1c,\n",
+    "    train_loader,\n",
+    "    val_loader,\n",
+    "    device,\n",
+    "    epochs=100,\n",
+    "    lr=1e-3,\n",
+    "    patience=20\n",
+    ")\n",
+    "scoreboard.add(result_2_1c)\n",
+    "\n",
+    "print(\"\\n[EXP 2.1c] Random Anchors for comparison...\")\n",
+    "model_2_1c_rand = HopfAnchorModel(num_anchors=64).to(device)\n",
+    "result_2_1c_rand = run_experiment(\n",
+    "    \"2.1c_RandomAnchors\",\n",
+    "    model_2_1c_rand,\n",
+    "    train_loader,\n",
+    "    val_loader,\n",
+    "    device,\n",
+    "    epochs=100,\n",
+    "    lr=1e-3,\n",
+    "    patience=20\n",
+    ")\n",
+    "scoreboard.add(result_2_1c_rand)\n",
+    ""
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Experiment 2.1d: Hierarchical Constellation"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "\n",
+    "Coarse anchors on base S\u00b2, fine anchors per fiber.\n",
+    "Two-level constellation system.\n",
+    ""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "\n",
+    "class HierarchicalConstellation(nn.Module):\n",
+    "    '''Two-level: coarse S\u00b2 + fine fiber anchors.'''\n",
+    "    def __init__(self, embed_dim, coarse_anchors=8, fine_per_fiber=8):\n",
+    "        super().__init__()\n",
+    "        self.embed_dim = embed_dim\n",
+    "        self.coarse_anchors = coarse_anchors\n",
+    "        self.fine_per_fiber = fine_per_fiber\n",
+    "\n",
+    "        # Coarse: S\u00b2 anchors\n",
+    "        coarse = F.normalize(torch.randn(coarse_anchors, 3), dim=-1)\n",
+    "        self.register_buffer(\"coarse\", coarse)\n",
+    "\n",
+    "        # Fine: per-coarse-anchor fiber codes\n",
+    "        fine = F.normalize(torch.randn(coarse_anchors * fine_per_fiber, 1), dim=-1)\n",
+    "        self.register_buffer(\"fine\", fine)\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        # x: (B, 4) on S\u00b3\n",
+    "        # Decompose as (base on S\u00b2, fiber)\n",
+    "        x_norm = l2_normalize(x)\n",
+    "        base_3d = x_norm[:, :3]\n",
+    "        fiber_1d = x_norm[:, 3:4]\n",
+    "\n",
+    "        # Coarse distances on S\u00b2\n",
+    "        cos_sim_coarse = torch.mm(base_3d, self.coarse.T)  # (B, coarse_anchors)\n",
+    "        dist_coarse = 1.0 - cos_sim_coarse\n",
+    "\n",
+    "        # Fine distances on fiber (per coarse anchor)\n",
+    "        cos_sim_fine = torch.mm(fiber_1d, self.fine.T)  # (B, coarse*fine)\n",
+    "        dist_fine = 1.0 - cos_sim_fine\n",
+    "\n",
+    "        # Concatenate coarse and fine\n",
+    "        distances = torch.cat([dist_coarse, dist_fine], dim=1)  # (B, coarse+coarse*fine)\n",
+    "        return distances\n",
+    "\n",
+    "class HierarchicalModel(nn.Module):\n",
+    "    def __init__(self):\n",
+    "        super().__init__()\n",
+    "        self.frontend = Scattering2DFeatures(out_dim=128)\n",
+    "        self.to_4d = nn.Linear(128, 4)\n",
+    "        self.constellation = HierarchicalConstellation(embed_dim=4, coarse_anchors=8, fine_per_fiber=8)\n",
+    "        total_anchors = 8 + 8*8\n",
+    "        self.patchwork = Patchwork(total_anchors, num_blocks=2)\n",
+    "        self.head = nn.Linear(total_anchors, 10)\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        feat = self.frontend(x)\n",
+    "        z_4d = self.to_4d(feat)\n",
+    "        z_norm = l2_normalize(z_4d)\n",
+    "        tri = self.constellation(z_norm)\n",
+    "        pw = self.patchwork(tri)\n",
+    "        logits = self.head(pw)\n",
+    "        return logits, z_norm\n",
+    "\n",
+    "print(\"\\n[EXP 2.1d] Hierarchical Constellation...\")\n",
+    "model_2_1d = HierarchicalModel().to(device)\n",
+    "result_2_1d = run_experiment(\n",
+    "    \"2.1d_HierarchicalConst\",\n",
+    "    model_2_1d,\n",
+    "    train_loader,\n",
+    "    val_loader,\n",
+    "    device,\n",
+    "    epochs=100,\n",
+    "    lr=1e-3,\n",
+    "    patience=20\n",
+    ")\n",
+    "scoreboard.add(result_2_1d)\n",
+    ""
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Experiment 2.2a: Grassmannian \u2014 Class Subspaces"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "\n",
+    "Each class = k-dim subspace of \u211d\u207f.\n",
+    "Distance via principal angles (SVD-based).\n",
+    "Literature: +1.3% over single vectors.\n",
+    ""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "\n",
+    "class GrassmannianClassifier(nn.Module):\n",
+    "    '''Each class is a k-dim subspace.'''\n",
+    "    def __init__(self, embed_dim=64, num_classes=10, k=5):\n",
+    "        super().__init__()\n",
+    "        self.embed_dim = embed_dim\n",
+    "        self.num_classes = num_classes\n",
+    "        self.k = k\n",
+    "\n",
+    "        # Each class: k orthonormal vectors\n",
+    "        U_init = []\n",
+    "        for _ in range(num_classes):\n",
+    "            u, _ = torch.linalg.qr(torch.randn(embed_dim, k))\n",
+    "            U_init.append(u)\n",
+    "        self.U = nn.ParameterList([nn.Parameter(u, requires_grad=True) for u in U_init])\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        # x: (B, embed_dim) on S^{d-1}\n",
+    "        x = l2_normalize(x)  # Normalize\n",
+    "        B = x.shape[0]\n",
+    "\n",
+    "        # Distance to each class subspace via principal angles\n",
+    "        distances = []\n",
+    "        for c in range(self.num_classes):\n",
+    "            U_c = self.U[c]  # (embed_dim, k)\n",
+    "            # Projection: ||x - U_c U_c^T x||\n",
+    "            proj = torch.mm(U_c, U_c.T)  # (embed_dim, embed_dim)\n",
+    "            projected = torch.mm(x, proj.T)  # (B, embed_dim)\n",
+    "            dist = torch.norm(x - projected, dim=1)  # (B,)\n",
+    "            distances.append(dist)\n",
+    "\n",
+    "        distances = torch.stack(distances, dim=1)  # (B, num_classes)\n",
+    "        logits = -distances  # Negative distance = similarity\n",
+    "        return logits, x\n",
+    "\n",
+    "class GrassmannianModel(nn.Module):\n",
+    "    def __init__(self):\n",
+    "        super().__init__()\n",
+    "        self.encoder = Scattering2DFeatures(out_dim=128)\n",
+    "        self.to_embedding = nn.Linear(128, 64)\n",
+    "        self.classifier = GrassmannianClassifier(embed_dim=64, num_classes=10, k=5)\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        feat = self.encoder(x)\n",
+    "        emb = self.to_embedding(feat)\n",
+    "        emb = l2_normalize(emb)\n",
+    "        logits, _ = self.classifier(emb)\n",
+    "        return logits, emb\n",
+    "\n",
+    "print(\"\\n[EXP 2.2a] Grassmannian Class Subspaces...\")\n",
+    "model_2_2a = GrassmannianModel().to(device)\n",
+    "result_2_2a = run_experiment(\n",
+    "    \"2.2a_GrassmannianClasses\",\n",
+    "    model_2_2a,\n",
+    "    train_loader,\n",
+    "    val_loader,\n",
+    "    device,\n",
+    "    epochs=100,\n",
+    "    lr=1e-3,\n",
+    "    patience=20\n",
+    ")\n",
+    "scoreboard.add(result_2_2a)\n",
+    ""
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Experiment 2.2b: Grassmannian Distance Metrics"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "\n",
+    "Compare:\n",
+    "- Geodesic distance on Grassmannian\n",
+    "- Chordal distance\n",
+    "- Projection distance (subspace angle)\n",
+    ""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "\n",
+    "class GrassmannianMetricAblation(nn.Module):\n",
+    "    '''Grassmannian with configurable metric.'''\n",
+    "    def __init__(self, embed_dim=64, num_classes=10, k=5, metric=\"geodesic\"):\n",
+    "        super().__init__()\n",
+    "        self.embed_dim = embed_dim\n",
+    "        self.num_classes = num_classes\n",
+    "        self.k = k\n",
+    "        self.metric = metric\n",
+    "\n",
+    "        U_init = []\n",
+    "        for _ in range(num_classes):\n",
+    "            u, _ = torch.linalg.qr(torch.randn(embed_dim, k))\n",
+    "            U_init.append(u)\n",
+    "        self.U = nn.ParameterList([nn.Parameter(u, requires_grad=True) for u in U_init])\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        x = l2_normalize(x)\n",
+    "        B = x.shape[0]\n",
+    "        distances = []\n",
+    "\n",
+    "        for c in range(self.num_classes):\n",
+    "            U_c = self.U[c]\n",
+    "\n",
+    "            if self.metric == \"projection\":\n",
+    "                # Subspace angle: arccos(||U_c^T x||)\n",
+    "                proj_norm = torch.norm(torch.mm(x, U_c), dim=1)\n",
+    "                dist = torch.acos(torch.clamp(proj_norm, -1+1e-6, 1-1e-6))\n",
+    "            elif self.metric == \"chordal\":\n",
+    "                # Chordal: ||U_c U_c^T - x x^T||_F\n",
+    "                proj_c = torch.mm(U_c, U_c.T)\n",
+    "                xx_t = torch.bmm(x.unsqueeze(2), x.unsqueeze(1))  # (B, embed_dim, embed_dim)\n",
+    "                dist = torch.norm(proj_c.unsqueeze(0) - xx_t, dim=(1,2))\n",
+    "            else:  # geodesic\n",
+    "                # Geodesic approximation via subspace angles\n",
+    "                proj_norm = torch.norm(torch.mm(x, U_c), dim=1)\n",
+    "                dist = torch.acos(torch.clamp(proj_norm, -1+1e-6, 1-1e-6)) ** 2\n",
+    "\n",
+    "            distances.append(dist)\n",
+    "\n",
+    "        distances = torch.stack(distances, dim=1)\n",
+    "        logits = -distances\n",
+    "        return logits, x\n",
+    "\n",
+    "class GrassmannianMetricModel(nn.Module):\n",
+    "    def __init__(self, metric=\"geodesic\"):\n",
+    "        super().__init__()\n",
+    "        self.encoder = Scattering2DFeatures(out_dim=128)\n",
+    "        self.to_embedding = nn.Linear(128, 64)\n",
+    "        self.classifier = GrassmannianMetricAblation(embed_dim=64, num_classes=10, k=5, metric=metric)\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        feat = self.encoder(x)\n",
+    "        emb = self.to_embedding(feat)\n",
+    "        emb = l2_normalize(emb)\n",
+    "        logits, _ = self.classifier(emb)\n",
+    "        return logits, emb\n",
+    "\n",
+    "# Test all three metrics\n",
+    "for metric in [\"projection\", \"chordal\", \"geodesic\"]:\n",
+    "    print(f\"\\n[EXP 2.2b] Testing {metric} metric...\")\n",
+    "    model = GrassmannianMetricModel(metric=metric).to(device)\n",
+    "    result = run_experiment(\n",
+    "        f\"2.2b_Grassmannian_{metric}\",\n",
+    "        model,\n",
+    "        train_loader,\n",
+    "        val_loader,\n",
+    "        device,\n",
+    "        epochs=100,\n",
+    "        lr=1e-3,\n",
+    "        patience=20\n",
+    "    )\n",
+    "    scoreboard.add(result)\n",
+    ""
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Experiment 2.2c: Per-Class Anchor Subspace"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "\n",
+    "Each anchor defines a subspace, not a point.\n",
+    "Extend constellation to subspace-valued anchors.\n",
+    ""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "\n",
+    "class SubspaceConstellation(nn.Module):\n",
+    "    '''Each anchor is a k-dim subspace.'''\n",
+    "    def __init__(self, embed_dim, num_anchors=16, k=3):\n",
+    "        super().__init__()\n",
+    "        self.embed_dim = embed_dim\n",
+    "        self.num_anchors = num_anchors\n",
+    "        self.k = k\n",
+    "\n",
+    "        # Each anchor: orthonormal basis\n",
+    "        U_list = []\n",
+    "        for _ in range(num_anchors):\n",
+    "            u, _ = torch.linalg.qr(torch.randn(embed_dim, k))\n",
+    "            U_list.append(u)\n",
+    "        self.U = nn.ParameterList([nn.Parameter(u, requires_grad=False) for u in U_list])\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        # x: (B, embed_dim)\n",
+    "        x = l2_normalize(x)\n",
+    "        B = x.shape[0]\n",
+    "        distances = []\n",
+    "\n",
+    "        for i in range(self.num_anchors):\n",
+    "            U_i = self.U[i]  # (embed_dim, k)\n",
+    "            # Distance: 1 - ||U_i^T x||\n",
+    "            proj_norm = torch.norm(torch.mm(x, U_i), dim=1)  # (B,)\n",
+    "            dist = 1.0 - proj_norm\n",
+    "            distances.append(dist)\n",
+    "\n",
+    "        distances = torch.stack(distances, dim=1)  # (B, num_anchors)\n",
+    "        return distances\n",
+    "\n",
+    "class SubspaceConstellationModel(nn.Module):\n",
+    "    def __init__(self):\n",
+    "        super().__init__()\n",
+    "        self.encoder = Scattering2DFeatures(out_dim=128)\n",
+    "        self.to_embedding = nn.Linear(128, 64)\n",
+    "        self.constellation = SubspaceConstellation(embed_dim=64, num_anchors=16, k=3)\n",
+    "        self.patchwork = Patchwork(16, num_blocks=2)\n",
+    "        self.head = nn.Linear(16, 10)\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        feat = self.encoder(x)\n",
+    "        emb = self.to_embedding(feat)\n",
+    "        emb = l2_normalize(emb)\n",
+    "        tri = self.constellation(emb)\n",
+    "        pw = self.patchwork(tri)\n",
+    "        logits = self.head(pw)\n",
+    "        return logits, emb\n",
+    "\n",
+    "print(\"\\n[EXP 2.2c] Per-Class Anchor Subspace...\")\n",
+    "model_2_2c = SubspaceConstellationModel().to(device)\n",
+    "result_2_2c = run_experiment(\n",
+    "    \"2.2c_SubspaceAnchors\",\n",
+    "    model_2_2c,\n",
+    "    train_loader,\n",
+    "    val_loader,\n",
+    "    device,\n",
+    "    epochs=100,\n",
+    "    lr=1e-3,\n",
+    "    patience=20\n",
+    ")\n",
+    "scoreboard.add(result_2_2c)\n",
+    ""
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Experiment 2.3a: Flag Manifold \u2014 Nested Hierarchy"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "\n",
+    "V\u2081 \u2282 V\u2082 \u2282 ... \u2282 V\u2096 nested subspaces.\n",
+    "Flag decomposition of frequency channels: DC \u2282 low \u2282 mid \u2282 high.\n",
+    ""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "\n",
+    "class FlagManifold(nn.Module):\n",
+    "    '''Nested subspace hierarchy on Flag(d; k1, k2, k3, d).'''\n",
+    "    def __init__(self, embed_dim=64, dims=[10, 20, 40, 64]):\n",
+    "        super().__init__()\n",
+    "        self.embed_dim = embed_dim\n",
+    "        self.dims = dims  # increasing sequence\n",
+    "\n",
+    "        # Build nested bases\n",
+    "        # V1 \u2282 V2 \u2282 ... \u2282 Vd\n",
+    "        bases = []\n",
+    "        A = torch.randn(embed_dim, embed_dim)\n",
+    "        Q, _ = torch.linalg.qr(A)\n",
+    "\n",
+    "        for d in dims:\n",
+    "            bases.append(Q[:, :d])\n",
+    "\n",
+    "        self.bases = nn.ParameterList([nn.Parameter(b, requires_grad=False) for b in bases])\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        # x: (B, embed_dim)\n",
+    "        x = l2_normalize(x)\n",
+    "        B = x.shape[0]\n",
+    "\n",
+    "        # Compute projections at each level\n",
+    "        features = []\n",
+    "        for base in self.bases:\n",
+    "            # Projection: ||base base^T x||\n",
+    "            proj_norm = torch.norm(torch.mm(x, base), dim=1)\n",
+    "            features.append(proj_norm.unsqueeze(-1))\n",
+    "\n",
+    "        features = torch.cat(features, dim=1)  # (B, len(dims))\n",
+    "        return features\n",
+    "\n",
+    "class FlagModel(nn.Module):\n",
+    "    def __init__(self):\n",
+    "        super().__init__()\n",
+    "        self.encoder = Scattering2DFeatures(out_dim=128)\n",
+    "        self.to_embedding = nn.Linear(128, 64)\n",
+    "        self.flag = FlagManifold(embed_dim=64, dims=[10, 20, 40, 64])\n",
+    "        self.patchwork = Patchwork(4, num_blocks=1)\n",
+    "        self.head = nn.Linear(4, 10)\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        feat = self.encoder(x)\n",
+    "        emb = self.to_embedding(feat)\n",
+    "        emb = l2_normalize(emb)\n",
+    "        flag_feat = self.flag(emb)\n",
+    "        pw = self.patchwork(flag_feat)\n",
+    "        logits = self.head(pw)\n",
+    "        return logits, emb\n",
+    "\n",
+    "print(\"\\n[EXP 2.3a] Flag Manifold \u2014 Nested Hierarchy...\")\n",
+    "model_2_3a = FlagModel().to(device)\n",
+    "result_2_3a = run_experiment(\n",
+    "    \"2.3a_FlagManifold\",\n",
+    "    model_2_3a,\n",
+    "    train_loader,\n",
+    "    val_loader,\n",
+    "    device,\n",
+    "    epochs=100,\n",
+    "    lr=1e-3,\n",
+    "    patience=20\n",
+    ")\n",
+    "scoreboard.add(result_2_3a)\n",
+    ""
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Experiment 2.3b: Flag-Structured Anchor Hierarchy"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "\n",
+    "Coarse-to-fine anchor hierarchy based on flag decomposition.\n",
+    ""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "\n",
+    "class FlagStructuredConstellation(nn.Module):\n",
+    "    '''Anchors organized in nested flag structure.'''\n",
+    "    def __init__(self, embed_dim=64, num_anchors=64):\n",
+    "        super().__init__()\n",
+    "        # Organize anchors into levels\n",
+    "        level_sizes = [8, 16, 24, 16]  # sum=64\n",
+    "        anchors = []\n",
+    "\n",
+    "        for level_size in level_sizes:\n",
+    "            level_anchors = F.normalize(torch.randn(level_size, embed_dim), dim=-1)\n",
+    "            anchors.append(level_anchors)\n",
+    "\n",
+    "        self.anchors = nn.ParameterList([nn.Parameter(a, requires_grad=False) for a in anchors])\n",
+    "        self.num_anchors = num_anchors\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        x = l2_normalize(x)\n",
+    "        distances = []\n",
+    "\n",
+    "        for level_anchors in self.anchors:\n",
+    "            cos_sim = torch.mm(x, level_anchors.T)\n",
+    "            dist = 1.0 - cos_sim\n",
+    "            distances.append(dist)\n",
+    "\n",
+    "        distances = torch.cat(distances, dim=1)  # (B, num_anchors)\n",
+    "        return distances\n",
+    "\n",
+    "class FlagAnchorModel(nn.Module):\n",
+    "    def __init__(self):\n",
+    "        super().__init__()\n",
+    "        self.encoder = Scattering2DFeatures(out_dim=128)\n",
+    "        self.to_embedding = nn.Linear(128, 64)\n",
+    "        self.constellation = FlagStructuredConstellation(embed_dim=64, num_anchors=64)\n",
+    "        self.patchwork = Patchwork(64, num_blocks=2)\n",
+    "        self.head = nn.Linear(64, 10)\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        feat = self.encoder(x)\n",
+    "        emb = self.to_embedding(feat)\n",
+    "        emb = l2_normalize(emb)\n",
+    "        tri = self.constellation(emb)\n",
+    "        pw = self.patchwork(tri)\n",
+    "        logits = self.head(pw)\n",
+    "        return logits, emb\n",
+    "\n",
+    "print(\"\\n[EXP 2.3b] Flag-Structured Anchors...\")\n",
+    "model_2_3b = FlagAnchorModel().to(device)\n",
+    "result_2_3b = run_experiment(\n",
+    "    \"2.3b_FlagAnchors\",\n",
+    "    model_2_3b,\n",
+    "    train_loader,\n",
+    "    val_loader,\n",
+    "    device,\n",
+    "    epochs=100,\n",
+    "    lr=1e-3,\n",
+    "    patience=20\n",
+    ")\n",
+    "scoreboard.add(result_2_3b)\n",
+    ""
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Experiment 2.4a: Von Mises-Fisher Soft Assignment"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "\n",
+    "Replace hard nearest-anchor with vMF soft posteriors.\n",
+    "p(j|x) = \u03b1_j f(x;\u03bc_j,\u03ba_j) / \u03a3 \u03b1_k f(x;\u03bc_k,\u03ba_k)\n",
+    "Learned \u03ba per anchor = adaptive influence radius.\n",
+    ""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "\n",
+    "class vMFSoftAssignment(nn.Module):\n",
+    "    '''Von Mises-Fisher soft assignment with learned concentration.'''\n",
+    "    def __init__(self, embed_dim, num_anchors=64):\n",
+    "        super().__init__()\n",
+    "        self.embed_dim = embed_dim\n",
+    "        self.num_anchors = num_anchors\n",
+    "\n",
+    "        anchors = F.normalize(torch.randn(num_anchors, embed_dim), dim=-1)\n",
+    "        self.anchors = nn.Parameter(anchors, requires_grad=False)\n",
+    "\n",
+    "        # Learned concentration per anchor\n",
+    "        self.kappa = nn.Parameter(torch.ones(num_anchors) * 10.0)\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        # x: (B, embed_dim) on S^{d-1}\n",
+    "        x = l2_normalize(x)\n",
+    "        B = x.shape[0]\n",
+    "\n",
+    "        # Cosine similarity\n",
+    "        cos_sim = torch.mm(x, self.anchors.T)  # (B, num_anchors)\n",
+    "\n",
+    "        # vMF probability: exp(\u03ba * cos_sim) (normalization absorbed in softmax)\n",
+    "        log_prob = self.kappa.unsqueeze(0) * cos_sim  # (B, num_anchors)\n",
+    "        probs = F.softmax(log_prob, dim=1)  # (B, num_anchors)\n",
+    "\n",
+    "        # Return soft assignment (can use as features)\n",
+    "        return probs, self.kappa\n",
+    "\n",
+    "class vMFModel(nn.Module):\n",
+    "    def __init__(self):\n",
+    "        super().__init__()\n",
+    "        self.encoder = Scattering2DFeatures(out_dim=128)\n",
+    "        self.to_embedding = nn.Linear(128, 64)\n",
+    "        self.vmf = vMFSoftAssignment(embed_dim=64, num_anchors=64)\n",
+    "        self.patchwork = Patchwork(64, num_blocks=2)\n",
+    "        self.head = nn.Linear(64, 10)\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        feat = self.encoder(x)\n",
+    "        emb = self.to_embedding(feat)\n",
+    "        emb = l2_normalize(emb)\n",
+    "        soft_assign, kappa = self.vmf(emb)\n",
+    "        pw = self.patchwork(soft_assign)\n",
+    "        logits = self.head(pw)\n",
+    "        return logits, emb\n",
+    "\n",
+    "print(\"\\n[EXP 2.4a] vMF Soft Assignment...\")\n",
+    "model_2_4a = vMFModel().to(device)\n",
+    "result_2_4a = run_experiment(\n",
+    "    \"2.4a_vMFSoftAssign\",\n",
+    "    model_2_4a,\n",
+    "    train_loader,\n",
+    "    val_loader,\n",
+    "    device,\n",
+    "    epochs=100,\n",
+    "    lr=1e-3,\n",
+    "    patience=20\n",
+    ")\n",
+    "scoreboard.add(result_2_4a)\n",
+    ""
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Experiment 2.4b: vMF Mixture EM for Anchor Init"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "\n",
+    "Use EM to initialize anchors instead of uniform hypersphere.\n",
+    "Pre-cluster training set embeddings with vMF mixture model.\n",
+    ""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "\n",
+    "class vMFMixtureEM:\n",
+    "    '''vMF mixture clustering via EM.'''\n",
+    "    def __init__(self, K=64, max_iter=20, tol=1e-4):\n",
+    "        self.K = K\n",
+    "        self.max_iter = max_iter\n",
+    "        self.tol = tol\n",
+    "\n",
+    "    def fit(self, X):\n",
+    "        '''X: (N, d) data on sphere.'''\n",
+    "        N, d = X.shape\n",
+    "        X = X / (np.linalg.norm(X, axis=1, keepdims=True) + 1e-8)\n",
+    "\n",
+    "        # Initialize\n",
+    "        mu = X[np.random.choice(N, self.K, replace=False)]  # (K, d)\n",
+    "        kappa = np.ones(self.K) * 5.0\n",
+    "        alpha = np.ones(self.K) / self.K\n",
+    "\n",
+    "        for it in range(self.max_iter):\n",
+    "            # E-step\n",
+    "            cos_sim = X @ mu.T  # (N, K)\n",
+    "            log_prob = np.log(alpha + 1e-8) + kappa * cos_sim\n",
+    "            log_prob -= log_prob.max(axis=1, keepdims=True)\n",
+    "            gamma = np.exp(log_prob)\n",
+    "            gamma /= (gamma.sum(axis=1, keepdims=True) + 1e-8)\n",
+    "\n",
+    "            # M-step\n",
+    "            Nk = gamma.sum(axis=0)\n",
+    "            alpha = Nk / N\n",
+    "\n",
+    "            for k in range(self.K):\n",
+    "                mu_k = (gamma[:, k:k+1] * X).sum(axis=0)\n",
+    "                mu_k = mu_k / (np.linalg.norm(mu_k) + 1e-8)\n",
+    "                mu[k] = mu_k\n",
+    "\n",
+    "                rbar = (gamma[:, k] * (X @ mu_k)).sum() / (Nk[k] + 1e-8)\n",
+    "                kappa[k] = (d * rbar - 1) / (1 - rbar**2 + 1e-8)\n",
+    "\n",
+    "        return mu, kappa, alpha\n",
+    "\n",
+    "class EMInitModel(nn.Module):\n",
+    "    def __init__(self):\n",
+    "        super().__init__()\n",
+    "        self.encoder = Scattering2DFeatures(out_dim=128)\n",
+    "        self.to_embedding = nn.Linear(128, 64)\n",
+    "\n",
+    "        # Placeholder constellation (will be initialized after EM)\n",
+    "        self.constellation = Constellation(embed_dim=64, num_anchors=64, init=\"uniform\")\n",
+    "        self.patchwork = Patchwork(64, num_blocks=2)\n",
+    "        self.head = nn.Linear(64, 10)\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        feat = self.encoder(x)\n",
+    "        emb = self.to_embedding(feat)\n",
+    "        emb = l2_normalize(emb)\n",
+    "        tri = self.constellation(emb)\n",
+    "        pw = self.patchwork(tri)\n",
+    "        logits = self.head(pw)\n",
+    "        return logits, emb\n",
+    "\n",
+    "print(\"\\n[EXP 2.4b] vMF Mixture EM Anchor Init...\")\n",
+    "\n",
+    "# Collect training embeddings\n",
+    "print(\"  Collecting training embeddings...\")\n",
+    "all_emb = []\n",
+    "all_labels = []\n",
+    "model_temp = EMInitModel().to(device)\n",
+    "with torch.no_grad():\n",
+    "    for x, y in train_loader:\n",
+    "        x = x.to(device)\n",
+    "        feat = model_temp.encoder(x)\n",
+    "        emb = model_temp.to_embedding(feat)\n",
+    "        emb = l2_normalize(emb)\n",
+    "        all_emb.append(emb.cpu().numpy())\n",
+    "        all_labels.append(y.numpy())\n",
+    "\n",
+    "all_emb = np.vstack(all_emb)\n",
+    "print(f\"  Embeddings shape: {all_emb.shape}\")\n",
+    "\n",
+    "# EM fit\n",
+    "em = vMFMixtureEM(K=64, max_iter=20)\n",
+    "mu_em, kappa_em, alpha_em = em.fit(all_emb)\n",
+    "print(f\"  EM fit complete. Mean \u03ba: {kappa_em.mean():.2f}\")\n",
+    "\n",
+    "# Initialize model with EM anchors\n",
+    "model_2_4b = EMInitModel().to(device)\n",
+    "with torch.no_grad():\n",
+    "    model_2_4b.constellation.anchors.copy_(torch.tensor(mu_em, dtype=torch.float32).to(device))\n",
+    "\n",
+    "result_2_4b = run_experiment(\n",
+    "    \"2.4b_vMFEM_Init\",\n",
+    "    model_2_4b,\n",
+    "    train_loader,\n",
+    "    val_loader,\n",
+    "    device,\n",
+    "    epochs=100,\n",
+    "    lr=1e-3,\n",
+    "    patience=20\n",
+    ")\n",
+    "scoreboard.add(result_2_4b)\n",
+    ""
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Experiment 2.4c: vMF \u03ba Diagnostics"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "\n",
+    "Track per-class \u03ba convergence over training.\n",
+    "Monitor adaptive influence radius evolution.\n",
+    ""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "\n",
+    "class vMFDiagnosticsModel(nn.Module):\n",
+    "    '''vMF with per-class concentration tracking.'''\n",
+    "    def __init__(self):\n",
+    "        super().__init__()\n",
+    "        self.encoder = Scattering2DFeatures(out_dim=128)\n",
+    "        self.to_embedding = nn.Linear(128, 64)\n",
+    "\n",
+    "        # Per-class anchors and kappas\n",
+    "        self.num_classes = 10\n",
+    "        anchors = []\n",
+    "        kappas = []\n",
+    "        for c in range(self.num_classes):\n",
+    "            a = F.normalize(torch.randn(8, 64), dim=-1)  # 8 anchors per class\n",
+    "            anchors.append(a)\n",
+    "            k = nn.Parameter(torch.ones(8) * 5.0)\n",
+    "            kappas.append(k)\n",
+    "\n",
+    "        self.anchors = nn.ParameterList([nn.Parameter(a, requires_grad=False) for a in anchors])\n",
+    "        self.kappas = nn.ParameterList(kappas)\n",
+    "\n",
+    "        self.patchwork = Patchwork(80, num_blocks=2)\n",
+    "        self.head = nn.Linear(80, 10)\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        feat = self.encoder(x)\n",
+    "        emb = self.to_embedding(feat)\n",
+    "        emb = l2_normalize(emb)\n",
+    "\n",
+    "        # Soft assignment per class\n",
+    "        all_probs = []\n",
+    "        for c in range(self.num_classes):\n",
+    "            cos_sim = torch.mm(emb, self.anchors[c].T)  # (B, 8)\n",
+    "            log_prob = self.kappas[c].unsqueeze(0) * cos_sim\n",
+    "            probs = F.softmax(log_prob, dim=1)\n",
+    "            all_probs.append(probs)\n",
+    "\n",
+    "        combined = torch.cat(all_probs, dim=1)  # (B, 80)\n",
+    "        pw = self.patchwork(combined)\n",
+    "        logits = self.head(pw)\n",
+    "        return logits, emb\n",
+    "\n",
+    "print(\"\\n[EXP 2.4c] vMF \u03ba Diagnostics...\")\n",
+    "model_2_4c = vMFDiagnosticsModel().to(device)\n",
+    "\n",
+    "# Custom training with \u03ba tracking\n",
+    "writer_2_4c = make_writer(\"2.4c_vMF_kappa_diagnostic\")\n",
+    "criterion = nn.CrossEntropyLoss()\n",
+    "optimizer = torch.optim.AdamW(model_2_4c.parameters(), lr=1e-3, weight_decay=1e-4)\n",
+    "\n",
+    "for epoch in range(1, 51):  # Shorter run for diagnostics\n",
+    "    model_2_4c.train()\n",
+    "    total_loss = 0\n",
+    "    correct = 0\n",
+    "    total = 0\n",
+    "\n",
+    "    for images, labels in train_loader:\n",
+    "        images, labels = images.to(device), labels.to(device)\n",
+    "        optimizer.zero_grad()\n",
+    "\n",
+    "        logits, _ = model_2_4c(images)\n",
+    "        loss = criterion(logits, labels)\n",
+    "        loss.backward()\n",
+    "        torch.nn.utils.clip_grad_norm_(model_2_4c.parameters(), 1.0)\n",
+    "        optimizer.step()\n",
+    "\n",
+    "        total_loss += loss.item()\n",
+    "        _, pred = logits.max(1)\n",
+    "        correct += pred.eq(labels).sum().item()\n",
+    "        total += labels.size(0)\n",
+    "\n",
+    "    # Log \u03ba per class\n",
+    "    with torch.no_grad():\n",
+    "        for c in range(10):\n",
+    "            mean_kappa = model_2_4c.kappas[c].mean().item()\n",
+    "            max_kappa = model_2_4c.kappas[c].max().item()\n",
+    "            writer_2_4c.add_scalar(f\"kappa/class_{c}_mean\", mean_kappa, epoch)\n",
+    "            writer_2_4c.add_scalar(f\"kappa/class_{c}_max\", max_kappa, epoch)\n",
+    "\n",
+    "    if epoch % 10 == 0:\n",
+    "        print(f\"  E{epoch} | Loss {total_loss/len(train_loader):.4f} | Acc {100*correct/total:.1f}%\")\n",
+    "\n",
+    "writer_2_4c.close()\n",
+    "print(\"[EXP 2.4c] Complete \u2014 \u03ba diagnostics logged to TensorBoard\")\n",
+    "\n",
+    "# Add a result entry\n",
+    "result_2_4c = {\"experiment\": \"2.4c_vMF_kappa_diag\", \"best_val_acc\": 0, \"best_epoch\": 50, \"params\": 0, \"final_train_acc\": 0}\n",
+    "scoreboard.add(result_2_4c)\n",
+    ""
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Experiment 2.5a: E\u2088 Lattice Anchors"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "\n",
+    "240 maximally separated points on S\u2077.\n",
+    "Construct E\u2088 root system vectors.\n",
+    "Superior packing efficiency over random.\n",
+    ""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "\n",
+    "def construct_e8_roots():\n",
+    "    '''Construct E\u2088 root system (240 roots).'''\n",
+    "    roots = []\n",
+    "\n",
+    "    # E\u2088 roots in 8D (simplified construction)\n",
+    "    # 1. Roots from roots of D\u2088 (\u00b11, \u00b11, 0, 0, 0, 0, 0, 0 permutations)\n",
+    "    for i in range(8):\n",
+    "        for j in range(i+1, 8):\n",
+    "            for si in [-1, 1]:\n",
+    "                for sj in [-1, 1]:\n",
+    "                    root = [0]*8\n",
+    "                    root[i] = si\n",
+    "                    root[j] = sj\n",
+    "                    roots.append(root)\n",
+    "\n",
+    "    # 2. Roots from E\u2088 additional structure\n",
+    "    # All coordinates \u00b11 with even number of minus signs\n",
+    "    for signs in np.ndindex(2**8):\n",
+    "        vec = np.array([-1 if s else 1 for s in signs], dtype=float)\n",
+    "        if sum(vec < 0) % 2 == 0:\n",
+    "            roots.append(vec.tolist())\n",
+    "\n",
+    "    roots = np.array(roots[:240], dtype=np.float32)\n",
+    "    roots = roots / (np.linalg.norm(roots, axis=1, keepdims=True) + 1e-8)\n",
+    "    return roots\n",
+    "\n",
+    "class E8ConstantsAnchorModel(nn.Module):\n",
+    "    def __init__(self):\n",
+    "        super().__init__()\n",
+    "        self.encoder = Scattering2DFeatures(out_dim=256)\n",
+    "        self.to_8d = nn.Linear(256, 8)\n",
+    "\n",
+    "        # E\u2088 roots as anchors\n",
+    "        e8_roots = construct_e8_roots()\n",
+    "        self.register_buffer(\"e8_anchors\", torch.tensor(e8_roots, dtype=torch.float32))\n",
+    "\n",
+    "        self.patchwork = Patchwork(240, num_blocks=3)\n",
+    "        self.head = nn.Linear(240, 10)\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        feat = self.encoder(x)\n",
+    "        emb = self.to_8d(feat)\n",
+    "        emb = l2_normalize(emb)\n",
+    "\n",
+    "        # Triangulation with E\u2088 anchors\n",
+    "        cos_sim = torch.mm(emb, self.e8_anchors.T)\n",
+    "        distances = 1.0 - cos_sim\n",
+    "\n",
+    "        pw = self.patchwork(distances)\n",
+    "        logits = self.head(pw)\n",
+    "        return logits, emb\n",
+    "\n",
+    "print(\"\\n[EXP 2.5a] E\u2088 Lattice Anchors (240 roots)...\")\n",
+    "model_2_5a = E8ConstantsAnchorModel().to(device)\n",
+    "result_2_5a = run_experiment(\n",
+    "    \"2.5a_E8Lattice\",\n",
+    "    model_2_5a,\n",
+    "    train_loader,\n",
+    "    val_loader,\n",
+    "    device,\n",
+    "    epochs=100,\n",
+    "    lr=1e-3,\n",
+    "    patience=20\n",
+    ")\n",
+    "scoreboard.add(result_2_5a)\n",
+    ""
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Experiment 2.5b: Spherical t-Design Init"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "\n",
+    "Compare vs uniform hypersphere initialization.\n",
+    "t-designs have optimal polynomial approximation properties.\n",
+    ""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "\n",
+    "def spherical_t_design_init(d, N, t=2):\n",
+    "    '''Generate spherical t-design points on S^{d-1}.'''\n",
+    "    # Simplified: use QR-based + repulsive refinement\n",
+    "    A = torch.randn(N, d)\n",
+    "    Q, _ = torch.linalg.qr(A.T)\n",
+    "    points = Q.T[:N]\n",
+    "\n",
+    "    if N > d:\n",
+    "        extra = F.normalize(torch.randn(N - d, d), dim=-1)\n",
+    "        points = torch.cat([points, extra], dim=0)\n",
+    "\n",
+    "    return F.normalize(points, dim=-1)\n",
+    "\n",
+    "class TDesignModel(nn.Module):\n",
+    "    def __init__(self):\n",
+    "        super().__init__()\n",
+    "        self.encoder = Scattering2DFeatures(out_dim=128)\n",
+    "        self.to_embedding = nn.Linear(128, 64)\n",
+    "\n",
+    "        # t-design anchors\n",
+    "        t_design = spherical_t_design_init(64, 64, t=2)\n",
+    "        self.register_buffer(\"t_design_anchors\", t_design)\n",
+    "\n",
+    "        self.patchwork = Patchwork(64, num_blocks=2)\n",
+    "        self.head = nn.Linear(64, 10)\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        feat = self.encoder(x)\n",
+    "        emb = self.to_embedding(feat)\n",
+    "        emb = l2_normalize(emb)\n",
+    "\n",
+    "        cos_sim = torch.mm(emb, self.t_design_anchors.T)\n",
+    "        distances = 1.0 - cos_sim\n",
+    "\n",
+    "        pw = self.patchwork(distances)\n",
+    "        logits = self.head(pw)\n",
+    "        return logits, emb\n",
+    "\n",
+    "# Compare: t-design vs uniform\n",
+    "for init_type, desc in [(\"tdesign\", \"Spherical t-Design\"), (\"uniform\", \"Uniform Random\")]:\n",
+    "    if init_type == \"uniform\":\n",
+    "        model = TDesignModel()\n",
+    "        model.t_design_anchors.copy_(F.normalize(torch.randn(64, 64), dim=-1))\n",
+    "    else:\n",
+    "        model = TDesignModel()\n",
+    "\n",
+    "    model = model.to(device)\n",
+    "    result = run_experiment(\n",
+    "        f\"2.5b_TDesign_{init_type}\",\n",
+    "        model,\n",
+    "        train_loader,\n",
+    "        val_loader,\n",
+    "        device,\n",
+    "        epochs=100,\n",
+    "        lr=1e-3,\n",
+    "        patience=20\n",
+    "    )\n",
+    "    scoreboard.add(result)\n",
+    ""
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Experiment 2.5c: Thomson Problem Solver"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "\n",
+    "Energy minimization for N anchors on S^d.\n",
+    "Compare: Coulomb energy vs QR + iterative repulsion.\n",
+    ""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "\n",
+    "def thomson_coulomb_energy(points):\n",
+    "    '''Compute Coulomb energy: sum of 1/||p_i - p_j|| for i != j.'''\n",
+    "    # points: (N, d)\n",
+    "    diff = points.unsqueeze(0) - points.unsqueeze(1)  # (N, N, d)\n",
+    "    dist = torch.norm(diff, dim=2)  # (N, N)\n",
+    "    dist = torch.where(dist > 0, dist, torch.ones_like(dist) * 1e10)  # Avoid division by 0\n",
+    "    energy = (1.0 / dist).sum() / 2  # Divide by 2 to avoid double counting\n",
+    "    return energy\n",
+    "\n",
+    "def solve_thomson_problem(d=64, N=64, steps=50):\n",
+    "    '''Gradient descent to minimize Coulomb energy.'''\n",
+    "    points = F.normalize(torch.randn(N, d), dim=-1)\n",
+    "    points.requires_grad = True\n",
+    "    optimizer = torch.optim.Adam([points], lr=0.01)\n",
+    "\n",
+    "    for step in range(steps):\n",
+    "        optimizer.zero_grad()\n",
+    "        points_norm = F.normalize(points, dim=-1)\n",
+    "        energy = thomson_coulomb_energy(points_norm)\n",
+    "        energy.backward()\n",
+    "        optimizer.step()\n",
+    "\n",
+    "        if (step + 1) % 10 == 0:\n",
+    "            with torch.no_grad():\n",
+    "                points_norm = F.normalize(points, dim=-1)\n",
+    "                energy = thomson_coulomb_energy(points_norm)\n",
+    "                print(f\"    Thomson step {step+1}/{steps}: E={energy:.4f}\")\n",
+    "\n",
+    "    return F.normalize(points.detach(), dim=-1)\n",
+    "\n",
+    "class ThomsonModel(nn.Module):\n",
+    "    def __init__(self, use_thomson=True):\n",
+    "        super().__init__()\n",
+    "        self.encoder = Scattering2DFeatures(out_dim=128)\n",
+    "        self.to_embedding = nn.Linear(128, 64)\n",
+    "\n",
+    "        if use_thomson:\n",
+    "            print(\"  Solving Thomson problem...\")\n",
+    "            anchors = solve_thomson_problem(d=64, N=64, steps=50)\n",
+    "        else:\n",
+    "            anchors = F.normalize(torch.randn(64, 64), dim=-1)\n",
+    "\n",
+    "        self.register_buffer(\"anchors\", anchors)\n",
+    "        self.patchwork = Patchwork(64, num_blocks=2)\n",
+    "        self.head = nn.Linear(64, 10)\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        feat = self.encoder(x)\n",
+    "        emb = self.to_embedding(feat)\n",
+    "        emb = l2_normalize(emb)\n",
+    "\n",
+    "        cos_sim = torch.mm(emb, self.anchors.T)\n",
+    "        distances = 1.0 - cos_sim\n",
+    "\n",
+    "        pw = self.patchwork(distances)\n",
+    "        logits = self.head(pw)\n",
+    "        return logits, emb\n",
+    "\n",
+    "print(\"\\n[EXP 2.5c] Thomson Problem Solver...\")\n",
+    "model_2_5c = ThomsonModel(use_thomson=True).to(device)\n",
+    "result_2_5c = run_experiment(\n",
+    "    \"2.5c_Thomson\",\n",
+    "    model_2_5c,\n",
+    "    train_loader,\n",
+    "    val_loader,\n",
+    "    device,\n",
+    "    epochs=100,\n",
+    "    lr=1e-3,\n",
+    "    patience=20\n",
+    ")\n",
+    "scoreboard.add(result_2_5c)\n",
+    ""
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Final Scoreboard \u2014 All Experiments"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "\n",
+    "print(\"\\n\" + \"=\"*80)\n",
+    "print(\"FINAL SCOREBOARD \u2014 CATEGORY 2 MANIFOLD STRUCTURES\")\n",
+    "print(\"=\"*80)\n",
+    "scoreboard._print_table()\n",
+    "\n",
+    "# Save to file\n",
+    "scoreboard.save(\"/sessions/optimistic-wizardly-dijkstra/mnt/outputs/experiment_2_scoreboard.json\")\n",
+    "\n",
+    "print(\"\\n[COMPLETE] All experiments completed!\")\n",
+    "print(\"[LOGS] Check TensorBoard with: %tensorboard --logdir=runs\")\n",
+    ""
+   ],
+   "execution_count": null,
+   "outputs": []
+  }
+ ]
+}

spectral/experiment_3_compact_representations.ipynb

@@ -0,0 +1,1547 @@
+{
+ "nbformat": 4,
+ "nbformat_minor": 5,
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.10.0"
+  },
+  "accelerator": "GPU",
+  "colab": {
+   "provenance": [],
+   "gpuType": "T4"
+  }
+ },
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "\n",
+    "# GeoLIP Hypersphere Experiments \u2014 Category 3: Compact Representations\n",
+    "\n",
+    "## Overview\n",
+    "We investigate **compact, low-dimensional representations** on hyperspheres.\n",
+    "Techniques:\n",
+    "- **Random Fourier Features (RFF)** \u2014 explicit kernel approximations\n",
+    "- **Johnson-Lindenstrauss (JL)** \u2014 near-isometric random projections\n",
+    "- **Compressed Sensing (CS)** \u2014 sparse measurement matrices\n",
+    "- **Fourier Positional Encoding** (Tancik/Mildenhall style)\n",
+    "- **Spherical Harmonics** \u2014 angular frequency decomposition on S^d\n",
+    "\n",
+    "Each maps to **Constellation** (K anchors) \u2192 **Triangulation Profile** (K-dim distances).\n",
+    "\n",
+    "**Hypothesis:** Careful projection design + hypersphere structure \u2192\n",
+    "near-perfect classification with 64-128d embeddings.\n",
+    ""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "\n",
+    "# \u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\n",
+    "# SHARED UTILITIES \u2014 GeoLIP Hypersphere Experiments\n",
+    "# \u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\n",
+    "import os, sys, time, math, json, datetime, warnings\n",
+    "import numpy as np\n",
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "import torch.nn.functional as F\n",
+    "from torch.utils.data import DataLoader\n",
+    "from torch.utils.tensorboard import SummaryWriter\n",
+    "import torchvision\n",
+    "import torchvision.transforms as T\n",
+    "import matplotlib\n",
+    "matplotlib.use(\"Agg\")\n",
+    "import matplotlib.pyplot as plt\n",
+    "from collections import defaultdict\n",
+    "\n",
+    "warnings.filterwarnings(\"ignore\", category=UserWarning)\n",
+    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
+    "print(f\"[DEVICE] {device}\" + (f\" \u2014 {torch.cuda.get_device_name()}\" if device.type == \"cuda\" else \"\"))\n",
+    "\n",
+    "# \u2500\u2500 CIFAR-10 Loaders \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
+    "CIFAR10_MEAN = (0.4914, 0.4822, 0.4465)\n",
+    "CIFAR10_STD  = (0.2470, 0.2435, 0.2616)\n",
+    "CIFAR10_CLASSES = (\"plane\",\"car\",\"bird\",\"cat\",\"deer\",\"dog\",\"frog\",\"horse\",\"ship\",\"truck\")\n",
+    "\n",
+    "def get_cifar10_loaders(batch_size=256, num_workers=2, augment=True):\n",
+    "    \"\"\"Return train_loader, val_loader for CIFAR-10.\"\"\"\n",
+    "    tf_train = T.Compose([\n",
+    "        T.RandomCrop(32, padding=4),\n",
+    "        T.RandomHorizontalFlip(),\n",
+    "        T.ToTensor(),\n",
+    "        T.Normalize(CIFAR10_MEAN, CIFAR10_STD),\n",
+    "    ]) if augment else T.Compose([T.ToTensor(), T.Normalize(CIFAR10_MEAN, CIFAR10_STD)])\n",
+    "    tf_val = T.Compose([T.ToTensor(), T.Normalize(CIFAR10_MEAN, CIFAR10_STD)])\n",
+    "\n",
+    "    train_ds = torchvision.datasets.CIFAR10(root=\"./data\", train=True,  download=True, transform=tf_train)\n",
+    "    val_ds   = torchvision.datasets.CIFAR10(root=\"./data\", train=False, download=True, transform=tf_val)\n",
+    "    train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True,  num_workers=num_workers, pin_memory=True, drop_last=True)\n",
+    "    val_loader   = DataLoader(val_ds,   batch_size=batch_size, shuffle=False, num_workers=num_workers, pin_memory=True)\n",
+    "    print(f\"[DATA] CIFAR-10: {len(train_ds)} train, {len(val_ds)} val, batch_size={batch_size}\")\n",
+    "    return train_loader, val_loader\n",
+    "\n",
+    "def get_cifar10_raw_loaders(batch_size=256, num_workers=2):\n",
+    "    \"\"\"Return loaders with only ToTensor (no normalization) for custom front-ends.\"\"\"\n",
+    "    tf = T.Compose([T.ToTensor()])\n",
+    "    train_ds = torchvision.datasets.CIFAR10(root=\"./data\", train=True,  download=True, transform=tf)\n",
+    "    val_ds   = torchvision.datasets.CIFAR10(root=\"./data\", train=False, download=True, transform=tf)\n",
+    "    train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True,  num_workers=num_workers, pin_memory=True, drop_last=True)\n",
+    "    val_loader   = DataLoader(val_ds,   batch_size=batch_size, shuffle=False, num_workers=num_workers, pin_memory=True)\n",
+    "    print(f\"[DATA] CIFAR-10 (raw): {len(train_ds)} train, {len(val_ds)} val\")\n",
+    "    return train_loader, val_loader\n",
+    "\n",
+    "# \u2500\u2500 L2 Normalization to Hypersphere \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
+    "def l2_normalize(x, dim=-1, eps=1e-8):\n",
+    "    \"\"\"Project to unit hypersphere S^{d-1}.\"\"\"\n",
+    "    return F.normalize(x, p=2, dim=dim, eps=eps)\n",
+    "\n",
+    "# \u2500\u2500 Constellation Module \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
+    "class Constellation(nn.Module):\n",
+    "    \"\"\"\n",
+    "    K anchors on S^{d-1}. Triangulation: compute distance from input to each anchor.\n",
+    "    Output: K-dim vector of (1 - cos_sim) distances = triangulation profile.\n",
+    "    \"\"\"\n",
+    "    def __init__(self, embed_dim, num_anchors=64, init=\"uniform\"):\n",
+    "        super().__init__()\n",
+    "        self.embed_dim = embed_dim\n",
+    "        self.num_anchors = num_anchors\n",
+    "        anchors = self._init_anchors(embed_dim, num_anchors, init)\n",
+    "        self.anchors = nn.Parameter(anchors, requires_grad=False)  # anchors are NOT learned\n",
+    "\n",
+    "    def _init_anchors(self, d, K, method):\n",
+    "        if method == \"uniform\":\n",
+    "            # QR-based uniform hypersphere initialization\n",
+    "            A = torch.randn(K, d)\n",
+    "            Q, _ = torch.linalg.qr(A.T)\n",
+    "            anchors = Q.T[:K]\n",
+    "            # If K > d, fill remaining with random unit vectors\n",
+    "            if K > d:\n",
+    "                extra = F.normalize(torch.randn(K - d, d), dim=-1)\n",
+    "                anchors = torch.cat([anchors, extra], dim=0)\n",
+    "            return F.normalize(anchors, dim=-1)\n",
+    "        elif method == \"random\":\n",
+    "            return F.normalize(torch.randn(K, d), dim=-1)\n",
+    "        else:\n",
+    "            return F.normalize(torch.randn(K, d), dim=-1)\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        \"\"\"x: (B, d) on S^{d-1} \u2192 (B, K) triangulation distances.\"\"\"\n",
+    "        x_norm = l2_normalize(x)\n",
+    "        cos_sim = torch.mm(x_norm, self.anchors.T)  # (B, K)\n",
+    "        distances = 1.0 - cos_sim  # in [0, 2]\n",
+    "        return distances\n",
+    "\n",
+    "    def push_anchors(self, embeddings, labels, num_classes=10, momentum=0.1):\n",
+    "        \"\"\"Push anchors toward class centroids (call periodically).\"\"\"\n",
+    "        with torch.no_grad():\n",
+    "            for c in range(num_classes):\n",
+    "                mask = (labels == c)\n",
+    "                if mask.sum() < 2:\n",
+    "                    continue\n",
+    "                centroid = l2_normalize(embeddings[mask].mean(dim=0, keepdim=True))\n",
+    "                # Assign anchors to this class based on nearest centroid\n",
+    "                start = c * (self.num_anchors // num_classes)\n",
+    "                end = start + (self.num_anchors // num_classes)\n",
+    "                self.anchors.data[start:end] = l2_normalize(\n",
+    "                    (1 - momentum) * self.anchors.data[start:end] + momentum * centroid\n",
+    "                )\n",
+    "\n",
+    "    def anchor_activation(self, x):\n",
+    "        \"\"\"Return fraction of anchors that are 'active' (nearest to at least one sample).\"\"\"\n",
+    "        x_norm = l2_normalize(x)\n",
+    "        cos_sim = torch.mm(x_norm, self.anchors.T)\n",
+    "        nearest = cos_sim.argmax(dim=1)\n",
+    "        active = len(nearest.unique())\n",
+    "        return active, self.num_anchors\n",
+    "\n",
+    "# \u2500\u2500 Patchwork Block \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
+    "class SquaredReLU(nn.Module):\n",
+    "    def forward(self, x):\n",
+    "        return F.relu(x) ** 2\n",
+    "\n",
+    "class PatchworkBlock(nn.Module):\n",
+    "    \"\"\"SquaredReLU + LayerNorm block for post-constellation processing.\"\"\"\n",
+    "    def __init__(self, dim):\n",
+    "        super().__init__()\n",
+    "        self.net = nn.Sequential(\n",
+    "            nn.Linear(dim, dim),\n",
+    "            SquaredReLU(),\n",
+    "            nn.LayerNorm(dim),\n",
+    "            nn.Linear(dim, dim),\n",
+    "            SquaredReLU(),\n",
+    "            nn.LayerNorm(dim),\n",
+    "        )\n",
+    "    def forward(self, x):\n",
+    "        return x + self.net(x)  # residual\n",
+    "\n",
+    "class Patchwork(nn.Module):\n",
+    "    \"\"\"Stack of PatchworkBlocks.\"\"\"\n",
+    "    def __init__(self, dim, num_blocks=2):\n",
+    "        super().__init__()\n",
+    "        self.blocks = nn.Sequential(*[PatchworkBlock(dim) for _ in range(num_blocks)])\n",
+    "    def forward(self, x):\n",
+    "        return self.blocks(x)\n",
+    "\n",
+    "# \u2500\u2500 Classifier Head \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
+    "class ConstellationClassifier(nn.Module):\n",
+    "    \"\"\"Full pipeline: encoder output \u2192 constellation \u2192 patchwork \u2192 classify.\"\"\"\n",
+    "    def __init__(self, embed_dim, num_anchors=64, num_classes=10, patchwork_blocks=2):\n",
+    "        super().__init__()\n",
+    "        self.constellation = Constellation(embed_dim, num_anchors)\n",
+    "        self.patchwork = Patchwork(num_anchors, patchwork_blocks)\n",
+    "        self.head = nn.Linear(num_anchors, num_classes)\n",
+    "\n",
+    "    def forward(self, embeddings):\n",
+    "        tri = self.constellation(embeddings)       # (B, K)\n",
+    "        pw  = self.patchwork(tri)                   # (B, K)\n",
+    "        logits = self.head(pw)                      # (B, 10)\n",
+    "        return logits, tri\n",
+    "\n",
+    "    def count_params(self):\n",
+    "        return sum(p.numel() for p in self.parameters() if p.requires_grad)\n",
+    "\n",
+    "# \u2500\u2500 Losses \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
+    "class CayleyMengerVolumeLoss(nn.Module):\n",
+    "    \"\"\"CV loss: encourage constellation volume in 0.18-0.25 band.\"\"\"\n",
+    "    def __init__(self, target_low=0.18, target_high=0.25):\n",
+    "        super().__init__()\n",
+    "        self.target_low = target_low\n",
+    "        self.target_high = target_high\n",
+    "\n",
+    "    def forward(self, tri):\n",
+    "        \"\"\"tri: (B, K) triangulation distances.\"\"\"\n",
+    "        # Approximate volume via determinant of distance submatrix\n",
+    "        B, K = tri.shape\n",
+    "        # Use std of triangulation as proxy for volume spread\n",
+    "        vol = tri.std(dim=1).mean()\n",
+    "        # Penalty if outside target band\n",
+    "        below = F.relu(self.target_low - vol)\n",
+    "        above = F.relu(vol - self.target_high)\n",
+    "        return below + above, vol.item()\n",
+    "\n",
+    "class InfoNCELoss(nn.Module):\n",
+    "    \"\"\"Self-supervised contrastive loss on augmented views.\"\"\"\n",
+    "    def __init__(self, temperature=0.07):\n",
+    "        super().__init__()\n",
+    "        self.temperature = temperature\n",
+    "\n",
+    "    def forward(self, z1, z2):\n",
+    "        \"\"\"z1, z2: (B, d) L2-normalized embeddings of two augmented views.\"\"\"\n",
+    "        z1 = l2_normalize(z1)\n",
+    "        z2 = l2_normalize(z2)\n",
+    "        B = z1.shape[0]\n",
+    "        z = torch.cat([z1, z2], dim=0)  # (2B, d)\n",
+    "        sim = torch.mm(z, z.T) / self.temperature  # (2B, 2B)\n",
+    "        # Mask out self-similarity\n",
+    "        mask = ~torch.eye(2*B, dtype=bool, device=z.device)\n",
+    "        sim = sim.masked_fill(~mask, -1e9)\n",
+    "        # Positive pairs: (i, i+B) and (i+B, i)\n",
+    "        labels = torch.cat([torch.arange(B, 2*B), torch.arange(B)]).to(z.device)\n",
+    "        loss = F.cross_entropy(sim, labels)\n",
+    "        return loss\n",
+    "\n",
+    "class ClassConditionalInfoNCE(nn.Module):\n",
+    "    \"\"\"Class-conditional contrastive: positive = same class, not same image.\"\"\"\n",
+    "    def __init__(self, temperature=0.07):\n",
+    "        super().__init__()\n",
+    "        self.temperature = temperature\n",
+    "\n",
+    "    def forward(self, z, labels):\n",
+    "        z = l2_normalize(z)\n",
+    "        B = z.shape[0]\n",
+    "        sim = torch.mm(z, z.T) / self.temperature\n",
+    "        # Mask: positive if same class, negative otherwise\n",
+    "        label_mask = (labels.unsqueeze(0) == labels.unsqueeze(1)).float()\n",
+    "        label_mask.fill_diagonal_(0)  # exclude self\n",
+    "        # For each sample, compute log-softmax over all non-self, weighted by positives\n",
+    "        logits = sim - torch.diag(torch.full((B,), 1e9, device=z.device))\n",
+    "        pos_mask = label_mask\n",
+    "        neg_mask = 1 - label_mask\n",
+    "        neg_mask.fill_diagonal_(0)\n",
+    "        # InfoNCE: -log(sum_pos exp(sim) / sum_all exp(sim))\n",
+    "        exp_sim = torch.exp(logits)\n",
+    "        pos_sum = (exp_sim * pos_mask).sum(dim=1)\n",
+    "        all_sum = (exp_sim * (pos_mask + neg_mask)).sum(dim=1)\n",
+    "        loss = -torch.log(pos_sum / (all_sum + 1e-8) + 1e-8).mean()\n",
+    "        return loss\n",
+    "\n",
+    "# \u2500\u2500 TensorBoard Helpers \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
+    "def make_writer(experiment_id, base_dir=\"runs\"):\n",
+    "    \"\"\"Create a TensorBoard SummaryWriter with timestamped run directory.\"\"\"\n",
+    "    timestamp = datetime.datetime.now().strftime(\"%Y%m%d_%H%M%S\")\n",
+    "    log_dir = os.path.join(base_dir, f\"{experiment_id}_{timestamp}\")\n",
+    "    writer = SummaryWriter(log_dir=log_dir)\n",
+    "    print(f\"[TB] Logging to {log_dir}\")\n",
+    "    return writer\n",
+    "\n",
+    "def log_epoch_metrics(writer, epoch, metrics_dict, prefix=\"\"):\n",
+    "    \"\"\"Log a dict of metrics to TensorBoard.\"\"\"\n",
+    "    for key, val in metrics_dict.items():\n",
+    "        tag = f\"{prefix}/{key}\" if prefix else key\n",
+    "        writer.add_scalar(tag, val, epoch)\n",
+    "\n",
+    "def log_embedding_viz(writer, embeddings, labels, epoch, tag=\"embeddings\"):\n",
+    "    \"\"\"Log embeddings for TensorBoard projector.\"\"\"\n",
+    "    meta = [CIFAR10_CLASSES[l] for l in labels]\n",
+    "    writer.add_embedding(embeddings, metadata=meta, global_step=epoch, tag=tag)\n",
+    "\n",
+    "def log_anchor_histogram(writer, constellation, epoch):\n",
+    "    \"\"\"Log anchor norm distribution.\"\"\"\n",
+    "    norms = constellation.anchors.data.norm(dim=1)\n",
+    "    writer.add_histogram(\"anchors/norms\", norms, epoch)\n",
+    "    cos_sim_matrix = torch.mm(constellation.anchors.data, constellation.anchors.data.T)\n",
+    "    off_diag = cos_sim_matrix[~torch.eye(cos_sim_matrix.shape[0], dtype=bool)].flatten()\n",
+    "    writer.add_histogram(\"anchors/inter_cosine_sim\", off_diag, epoch)\n",
+    "\n",
+    "def log_confusion_matrix(writer, all_preds, all_labels, epoch):\n",
+    "    \"\"\"Log confusion matrix as image to TensorBoard.\"\"\"\n",
+    "    from sklearn.metrics import confusion_matrix\n",
+    "    cm = confusion_matrix(all_labels, all_preds, labels=range(10))\n",
+    "    fig, ax = plt.subplots(1, 1, figsize=(8, 8))\n",
+    "    im = ax.imshow(cm, interpolation=\"nearest\", cmap=plt.cm.Blues)\n",
+    "    ax.set_title(f\"Confusion Matrix \u2014 Epoch {epoch}\")\n",
+    "    ax.set_xlabel(\"Predicted\")\n",
+    "    ax.set_ylabel(\"True\")\n",
+    "    tick_marks = range(10)\n",
+    "    ax.set_xticks(tick_marks)\n",
+    "    ax.set_xticklabels(CIFAR10_CLASSES, rotation=45, ha=\"right\")\n",
+    "    ax.set_yticks(tick_marks)\n",
+    "    ax.set_yticklabels(CIFAR10_CLASSES)\n",
+    "    for i in range(10):\n",
+    "        for j in range(10):\n",
+    "            ax.text(j, i, str(cm[i, j]), ha=\"center\", va=\"center\",\n",
+    "                    color=\"white\" if cm[i, j] > cm.max() / 2 else \"black\", fontsize=7)\n",
+    "    fig.tight_layout()\n",
+    "    writer.add_figure(\"confusion_matrix\", fig, epoch)\n",
+    "    plt.close(fig)\n",
+    "\n",
+    "def log_triangulation_stats(writer, tri, epoch):\n",
+    "    \"\"\"Log triangulation distance statistics.\"\"\"\n",
+    "    writer.add_scalar(\"tri/mean\", tri.mean().item(), epoch)\n",
+    "    writer.add_scalar(\"tri/std\", tri.std().item(), epoch)\n",
+    "    writer.add_scalar(\"tri/min\", tri.min().item(), epoch)\n",
+    "    writer.add_scalar(\"tri/max\", tri.max().item(), epoch)\n",
+    "    writer.add_histogram(\"tri/distances\", tri.detach().cpu(), epoch)\n",
+    "\n",
+    "def log_scatter_2d(writer, embeddings, labels, epoch, tag=\"scatter_2d\", method=\"pca\"):\n",
+    "    \"\"\"Log 2D scatter plot of embeddings projected via PCA or t-SNE.\"\"\"\n",
+    "    emb_np = embeddings.detach().cpu().numpy()\n",
+    "    lab_np = labels.detach().cpu().numpy() if torch.is_tensor(labels) else np.array(labels)\n",
+    "    if method == \"pca\":\n",
+    "        from sklearn.decomposition import PCA\n",
+    "        proj = PCA(n_components=2).fit_transform(emb_np)\n",
+    "    else:\n",
+    "        from sklearn.manifold import TSNE\n",
+    "        proj = TSNE(n_components=2, perplexity=30, n_iter=500).fit_transform(emb_np)\n",
+    "    fig, ax = plt.subplots(1, 1, figsize=(8, 8))\n",
+    "    for c in range(10):\n",
+    "        mask = lab_np == c\n",
+    "        ax.scatter(proj[mask, 0], proj[mask, 1], s=3, alpha=0.5, label=CIFAR10_CLASSES[c])\n",
+    "    ax.legend(markerscale=5, fontsize=8)\n",
+    "    ax.set_title(f\"{tag} ({method.upper()}) \u2014 Epoch {epoch}\")\n",
+    "    fig.tight_layout()\n",
+    "    writer.add_figure(f\"{tag}/{method}\", fig, epoch)\n",
+    "    plt.close(fig)\n",
+    "\n",
+    "# \u2500\u2500 Training Utilities \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
+    "def train_one_epoch(model, loader, optimizer, criterion, device, epoch=0,\n",
+    "                    scheduler=None, extra_losses=None):\n",
+    "    \"\"\"\n",
+    "    Generic training loop for one epoch.\n",
+    "    model: callable(batch) \u2192 (logits, embeddings_on_sphere)\n",
+    "    extra_losses: list of (name, loss_fn, weight) where loss_fn(embeddings) \u2192 scalar\n",
+    "    \"\"\"\n",
+    "    model.train()\n",
+    "    total_loss = 0\n",
+    "    total_ce = 0\n",
+    "    correct = 0\n",
+    "    total = 0\n",
+    "    loss_components = defaultdict(float)\n",
+    "\n",
+    "    for batch_idx, (images, labels) in enumerate(loader):\n",
+    "        images, labels = images.to(device), labels.to(device)\n",
+    "        optimizer.zero_grad()\n",
+    "\n",
+    "        logits, embeddings = model(images)\n",
+    "        ce_loss = criterion(logits, labels)\n",
+    "        loss = ce_loss\n",
+    "        loss_components[\"CE\"] += ce_loss.item()\n",
+    "\n",
+    "        if extra_losses:\n",
+    "            for name, loss_fn, weight in extra_losses:\n",
+    "                if name == \"CV\":\n",
+    "                    el, vol_val = loss_fn(embeddings)\n",
+    "                    loss_components[\"CV_vol\"] += vol_val\n",
+    "                elif name == \"InfoNCE_class\":\n",
+    "                    el = loss_fn(embeddings, labels)\n",
+    "                else:\n",
+    "                    el = loss_fn(embeddings)\n",
+    "                loss = loss + weight * el\n",
+    "                loss_components[name] += el.item()\n",
+    "\n",
+    "        loss.backward()\n",
+    "        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)\n",
+    "        optimizer.step()\n",
+    "        if scheduler:\n",
+    "            scheduler.step()\n",
+    "\n",
+    "        total_loss += loss.item()\n",
+    "        _, predicted = logits.max(1)\n",
+    "        total += labels.size(0)\n",
+    "        correct += predicted.eq(labels).sum().item()\n",
+    "\n",
+    "    n_batches = len(loader)\n",
+    "    metrics = {\n",
+    "        \"loss/total\": total_loss / n_batches,\n",
+    "        \"loss/CE\": loss_components[\"CE\"] / n_batches,\n",
+    "        \"accuracy/train\": 100.0 * correct / total,\n",
+    "    }\n",
+    "    for k, v in loss_components.items():\n",
+    "        if k not in (\"CE\",):\n",
+    "            metrics[f\"loss/{k}\"] = v / n_batches\n",
+    "    return metrics\n",
+    "\n",
+    "@torch.no_grad()\n",
+    "def evaluate(model, loader, criterion, device, collect_embeddings=False):\n",
+    "    \"\"\"Evaluate model on validation set. Optionally collect embeddings.\"\"\"\n",
+    "    model.eval()\n",
+    "    total_loss = 0\n",
+    "    correct = 0\n",
+    "    total = 0\n",
+    "    all_preds = []\n",
+    "    all_labels = []\n",
+    "    all_embeddings = []\n",
+    "\n",
+    "    for images, labels in loader:\n",
+    "        images, labels = images.to(device), labels.to(device)\n",
+    "        logits, embeddings = model(images)\n",
+    "        loss = criterion(logits, labels)\n",
+    "        total_loss += loss.item()\n",
+    "        _, predicted = logits.max(1)\n",
+    "        total += labels.size(0)\n",
+    "        correct += predicted.eq(labels).sum().item()\n",
+    "        all_preds.extend(predicted.cpu().numpy())\n",
+    "        all_labels.extend(labels.cpu().numpy())\n",
+    "        if collect_embeddings:\n",
+    "            all_embeddings.append(embeddings.cpu())\n",
+    "\n",
+    "    n_batches = len(loader)\n",
+    "    metrics = {\n",
+    "        \"loss/val\": total_loss / n_batches,\n",
+    "        \"accuracy/val\": 100.0 * correct / total,\n",
+    "    }\n",
+    "    result = {\"metrics\": metrics, \"preds\": np.array(all_preds), \"labels\": np.array(all_labels)}\n",
+    "    if collect_embeddings:\n",
+    "        result[\"embeddings\"] = torch.cat(all_embeddings, dim=0)\n",
+    "    return result\n",
+    "\n",
+    "def run_experiment(experiment_id, model, train_loader, val_loader, device,\n",
+    "                   epochs=100, lr=1e-3, weight_decay=1e-4, extra_losses=None,\n",
+    "                   push_anchors_every=50, log_embeddings_every=25,\n",
+    "                   patience=20):\n",
+    "    \"\"\"\n",
+    "    Full experiment runner with TensorBoard logging, early stopping, and diagnostics.\n",
+    "    model must return (logits, sphere_embeddings).\n",
+    "    \"\"\"\n",
+    "    writer = make_writer(experiment_id)\n",
+    "    criterion = nn.CrossEntropyLoss()\n",
+    "    optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=weight_decay)\n",
+    "    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs)\n",
+    "\n",
+    "    best_val_acc = 0.0\n",
+    "    best_epoch = 0\n",
+    "    no_improve = 0\n",
+    "\n",
+    "    # Log model info\n",
+    "    n_params = sum(p.numel() for p in model.parameters() if p.requires_grad)\n",
+    "    writer.add_text(\"config/params\", f\"Learned params: {n_params:,}\")\n",
+    "    writer.add_text(\"config/experiment\", experiment_id)\n",
+    "    print(f\"\\n{'='*70}\")\n",
+    "    print(f\"[EXP] {experiment_id}\")\n",
+    "    print(f\"[EXP] Learned params: {n_params:,}, Epochs: {epochs}, LR: {lr}\")\n",
+    "    print(f\"{'='*70}\")\n",
+    "\n",
+    "    for epoch in range(1, epochs + 1):\n",
+    "        t0 = time.time()\n",
+    "\n",
+    "        # Train\n",
+    "        train_metrics = train_one_epoch(\n",
+    "            model, train_loader, optimizer, criterion, device, epoch,\n",
+    "            extra_losses=extra_losses\n",
+    "        )\n",
+    "\n",
+    "        # Evaluate\n",
+    "        collect_emb = (epoch % log_embeddings_every == 0) or (epoch == epochs)\n",
+    "        val_result = evaluate(model, val_loader, criterion, device, collect_embeddings=collect_emb)\n",
+    "        val_metrics = val_result[\"metrics\"]\n",
+    "\n",
+    "        # Combine metrics\n",
+    "        all_metrics = {**train_metrics, **val_metrics}\n",
+    "        all_metrics[\"lr\"] = optimizer.param_groups[0][\"lr\"]\n",
+    "        all_metrics[\"epoch_time\"] = time.time() - t0\n",
+    "        all_metrics[\"train_val_gap\"] = train_metrics[\"accuracy/train\"] - val_metrics[\"accuracy/val\"]\n",
+    "\n",
+    "        # Constellation diagnostics\n",
+    "        if hasattr(model, \"constellation\"):\n",
+    "            constellation = model.constellation\n",
+    "        elif hasattr(model, \"classifier\") and hasattr(model.classifier, \"constellation\"):\n",
+    "            constellation = model.classifier.constellation\n",
+    "        else:\n",
+    "            constellation = None\n",
+    "\n",
+    "        if constellation is not None:\n",
+    "            log_anchor_histogram(writer, constellation, epoch)\n",
+    "\n",
+    "            # Push anchors periodically\n",
+    "            if epoch % push_anchors_every == 0 and \"embeddings\" in val_result:\n",
+    "                emb = val_result[\"embeddings\"][:2000].to(device)\n",
+    "                lab = torch.tensor(val_result[\"labels\"][:2000]).to(device)\n",
+    "                constellation.push_anchors(emb, lab)\n",
+    "                active, total_anch = constellation.anchor_activation(emb)\n",
+    "                all_metrics[\"anchors/active\"] = active\n",
+    "                all_metrics[\"anchors/total\"] = total_anch\n",
+    "                all_metrics[\"anchors/utilization\"] = active / total_anch\n",
+    "                print(f\"  [ANCHOR PUSH] Active: {active}/{total_anch} ({100*active/total_anch:.1f}%)\")\n",
+    "\n",
+    "        # Log all metrics\n",
+    "        log_epoch_metrics(writer, epoch, all_metrics)\n",
+    "        scheduler.step()\n",
+    "\n",
+    "        # Log embeddings visualization\n",
+    "        if collect_emb and \"embeddings\" in val_result:\n",
+    "            emb = val_result[\"embeddings\"][:2000]\n",
+    "            lab = val_result[\"labels\"][:2000]\n",
+    "            log_scatter_2d(writer, emb, lab, epoch, tag=\"val_embeddings\", method=\"pca\")\n",
+    "            if epoch == epochs:\n",
+    "                log_confusion_matrix(writer, val_result[\"preds\"], val_result[\"labels\"], epoch)\n",
+    "                log_scatter_2d(writer, emb, lab, epoch, tag=\"val_embeddings_tsne\", method=\"tsne\")\n",
+    "\n",
+    "        # Log triangulation stats from a sample batch\n",
+    "        if constellation is not None and collect_emb and \"embeddings\" in val_result:\n",
+    "            sample_emb = val_result[\"embeddings\"][:500].to(device)\n",
+    "            tri_sample = constellation(sample_emb)\n",
+    "            log_triangulation_stats(writer, tri_sample, epoch)\n",
+    "\n",
+    "        # Best model tracking\n",
+    "        val_acc = val_metrics[\"accuracy/val\"]\n",
+    "        if val_acc > best_val_acc:\n",
+    "            best_val_acc = val_acc\n",
+    "            best_epoch = epoch\n",
+    "            no_improve = 0\n",
+    "            writer.add_scalar(\"best/val_acc\", best_val_acc, epoch)\n",
+    "            writer.add_scalar(\"best/epoch\", best_epoch, epoch)\n",
+    "        else:\n",
+    "            no_improve += 1\n",
+    "\n",
+    "        # Print progress\n",
+    "        if epoch % 5 == 0 or epoch == 1 or epoch == epochs:\n",
+    "            print(f\"  E{epoch:3d} | TrainAcc {train_metrics['accuracy/train']:5.1f}% | \"\n",
+    "                  f\"ValAcc {val_acc:5.1f}% | Loss {all_metrics['loss/total']:.4f} | \"\n",
+    "                  f\"Best {best_val_acc:.1f}% @E{best_epoch} | \"\n",
+    "                  f\"Gap {all_metrics['train_val_gap']:+.1f}% | \"\n",
+    "                  f\"{all_metrics['epoch_time']:.1f}s\")\n",
+    "\n",
+    "        # Early stopping\n",
+    "        if no_improve >= patience:\n",
+    "            print(f\"  [EARLY STOP] No improvement for {patience} epochs. Best: {best_val_acc:.1f}% @E{best_epoch}\")\n",
+    "            break\n",
+    "\n",
+    "    writer.close()\n",
+    "    print(f\"\\n[RESULT] {experiment_id}: Best Val Acc = {best_val_acc:.2f}% at epoch {best_epoch}\")\n",
+    "    print(f\"[RESULT] Learned params: {n_params:,}\")\n",
+    "    return {\"experiment\": experiment_id, \"best_val_acc\": best_val_acc, \"best_epoch\": best_epoch,\n",
+    "            \"params\": n_params, \"final_train_acc\": train_metrics[\"accuracy/train\"]}\n",
+    "\n",
+    "# \u2500\u2500 Johnson-Lindenstrauss Projection \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
+    "class JLProjection(nn.Module):\n",
+    "    \"\"\"Fixed random Gaussian projection (not learned).\"\"\"\n",
+    "    def __init__(self, input_dim, output_dim):\n",
+    "        super().__init__()\n",
+    "        self.register_buffer(\"W\", torch.randn(input_dim, output_dim) / math.sqrt(output_dim))\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        return x @ self.W\n",
+    "\n",
+    "# \u2500\u2500 Results Scoreboard \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
+    "class Scoreboard:\n",
+    "    \"\"\"Track results across experiments in a session.\"\"\"\n",
+    "    def __init__(self):\n",
+    "        self.results = []\n",
+    "\n",
+    "    def add(self, result_dict):\n",
+    "        self.results.append(result_dict)\n",
+    "        self._print_table()\n",
+    "\n",
+    "    def _print_table(self):\n",
+    "        print(f\"\\n{'='*80}\")\n",
+    "        print(f\"{'SCOREBOARD':^80}\")\n",
+    "        print(f\"{'='*80}\")\n",
+    "        print(f\"{'Experiment':<40} {'Val Acc':>8} {'Params':>10} {'Best E':>7}\")\n",
+    "        print(f\"{'-'*40} {'-'*8} {'-'*10} {'-'*7}\")\n",
+    "        for r in sorted(self.results, key=lambda x: x[\"best_val_acc\"], reverse=True):\n",
+    "            print(f\"{r['experiment']:<40} {r['best_val_acc']:>7.2f}% {r['params']:>10,} {r['best_epoch']:>7}\")\n",
+    "        print(f\"{'='*80}\\n\")\n",
+    "\n",
+    "    def save(self, path=\"scoreboard.json\"):\n",
+    "        with open(path, \"w\") as f:\n",
+    "            json.dump(self.results, f, indent=2)\n",
+    "        print(f\"[SCOREBOARD] Saved to {path}\")\n",
+    "\n",
+    "scoreboard = Scoreboard()\n",
+    "print(\"[SHARED] All utilities loaded successfully.\")\n",
+    ""
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "\n",
+    "# Load CIFAR-10 (raw pixels for RFF on raw)\n",
+    "train_loader_raw, val_loader_raw = get_cifar10_raw_loaders(batch_size=256, num_workers=4)\n",
+    "\n",
+    "# Also load normalized for later experiments\n",
+    "train_loader, val_loader = get_cifar10_loaders(batch_size=256, num_workers=4, augment=True)\n",
+    ""
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "\n",
+    "## 3.1: Random Fourier Features (RFF)\n",
+    "\n",
+    "Random Fourier Features approximate RBF kernels explicitly:\n",
+    "$$z(x) = \\sqrt{\\frac{2}{D}} [\\cos(\\omega_1^T x + b_1), \\ldots, \\cos(\\omega_D^T x + b_D)]^T$$\n",
+    "\n",
+    "Key insight: RFF outputs are **on a hypersphere by construction**.\n",
+    ""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "\n",
+    "import torch.nn as nn\n",
+    "import torch.nn.functional as F\n",
+    "\n",
+    "class RFFLayer(nn.Module):\n",
+    "    \"\"\"Random Fourier Features layer.\"\"\"\n",
+    "    def __init__(self, input_dim, output_dim, scale=0.1):\n",
+    "        super().__init__()\n",
+    "        # Sample random frequencies\n",
+    "        self.register_buffer(\"W\", torch.randn(input_dim, output_dim) * scale)\n",
+    "        self.register_buffer(\"b\", torch.uniform_(torch.empty(output_dim), 0, 2 * math.pi))\n",
+    "        self.output_dim = output_dim\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        # x: (B, input_dim)\n",
+    "        # Compute cos(Wx + b)\n",
+    "        proj = torch.matmul(x, self.W) + self.b.unsqueeze(0)  # (B, output_dim)\n",
+    "        features = torch.cos(proj) * math.sqrt(2.0 / self.output_dim)  # (B, output_dim)\n",
+    "        # RFF is already on sphere by construction (approximately normalized)\n",
+    "        features = l2_normalize(features)\n",
+    "        return features\n",
+    "\n",
+    "class RFFPipeline(nn.Module):\n",
+    "    \"\"\"RFF on raw pixels \u2192 L2 normalize \u2192 Constellation classifier.\"\"\"\n",
+    "    def __init__(self, input_dim=3072, rff_dim=128, num_anchors=64, num_classes=10):\n",
+    "        super().__init__()\n",
+    "        self.rff = RFFLayer(input_dim, rff_dim, scale=0.1)\n",
+    "        self.classifier = ConstellationClassifier(rff_dim, num_anchors, num_classes)\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        # Flatten\n",
+    "        B = x.shape[0]\n",
+    "        x_flat = x.reshape(B, -1)  # (B, 3072)\n",
+    "        # RFF\n",
+    "        z = self.rff(x_flat)  # (B, 128) on sphere\n",
+    "        # Constellation + classification\n",
+    "        logits, tri = self.classifier(z)\n",
+    "        return logits, z\n",
+    "\n",
+    "# Experiment 3.1a\n",
+    "model_3_1a = RFFPipeline(input_dim=3072, rff_dim=128, num_anchors=64, num_classes=10).to(device)\n",
+    "result_3_1a = run_experiment(\n",
+    "    \"3.1a_RFF_RawPixels\",\n",
+    "    model_3_1a,\n",
+    "    train_loader_raw,\n",
+    "    val_loader_raw,\n",
+    "    device,\n",
+    "    epochs=100,\n",
+    "    lr=1e-3,\n",
+    "    weight_decay=1e-4,\n",
+    "    patience=20\n",
+    ")\n",
+    "scoreboard.add(result_3_1a)\n",
+    ""
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "\n",
+    "from kymatio.torch import Scattering2D\n",
+    "\n",
+    "class RFFScatteringPipeline(nn.Module):\n",
+    "    \"\"\"Scattering2D \u2192 flatten \u2192 RFF \u2192 constellation.\"\"\"\n",
+    "    def __init__(self, J=2, rff_dim=128, num_anchors=64, num_classes=10):\n",
+    "        super().__init__()\n",
+    "        self.scattering = Scattering2D(J=J, shape=(32, 32))\n",
+    "        # Scattering output size: 1 + 2*J (L=1) + J^2 (L=2) = 1 + 4 + 4 = 9 scales\n",
+    "        # Per pixel: 3 input channels \u2192 multiple scales\n",
+    "        # Empirically ~650 coefficients for 32x32\n",
+    "        self.scatter_dim = None  # Will compute on first forward pass\n",
+    "        self.rff = None\n",
+    "        self.rff_dim = rff_dim\n",
+    "        self.num_anchors = num_anchors\n",
+    "        self.num_classes = num_classes\n",
+    "        self.classifier = None\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        # x: (B, 3, 32, 32)\n",
+    "        # Scattering\n",
+    "        B = x.shape[0]\n",
+    "        # Scattering2D expects (B, C, H, W)\n",
+    "        scatter_out = self.scattering(x)  # (B, C*scales, H', W')\n",
+    "        # Flatten\n",
+    "        scatter_flat = scatter_out.reshape(B, -1)  # (B, scatter_dim)\n",
+    "\n",
+    "        # Initialize RFF on first pass\n",
+    "        if self.rff is None:\n",
+    "            self.scatter_dim = scatter_flat.shape[1]\n",
+    "            self.rff = RFFLayer(self.scatter_dim, self.rff_dim, scale=0.05).to(x.device)\n",
+    "            self.classifier = ConstellationClassifier(self.rff_dim, self.num_anchors, self.num_classes).to(x.device)\n",
+    "\n",
+    "        # RFF\n",
+    "        z = self.rff(scatter_flat)  # (B, rff_dim) on sphere\n",
+    "\n",
+    "        # Classification\n",
+    "        logits, tri = self.classifier(z)\n",
+    "        return logits, z\n",
+    "\n",
+    "model_3_1b = RFFScatteringPipeline(J=2, rff_dim=128, num_anchors=64, num_classes=10).to(device)\n",
+    "result_3_1b = run_experiment(\n",
+    "    \"3.1b_RFF_Scattering\",\n",
+    "    model_3_1b,\n",
+    "    train_loader_raw,\n",
+    "    val_loader_raw,\n",
+    "    device,\n",
+    "    epochs=100,\n",
+    "    lr=1e-3,\n",
+    "    weight_decay=1e-4,\n",
+    "    patience=20\n",
+    ")\n",
+    "scoreboard.add(result_3_1b)\n",
+    ""
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "\n",
+    "class FourierPositionalEncoding(nn.Module):\n",
+    "    \"\"\"\n",
+    "    Tancik/Mildenhall style positional encoding:\n",
+    "    \u03b3(v) = [cos(2\u03c0*B\u2080*v), sin(2\u03c0*B\u2080*v), cos(2\u03c0*B\u2081*v), sin(2\u03c0*B\u2081*v), ...]\n",
+    "    where B_i = 2^i for i=0..L-1\n",
+    "    Directly projects to hypersphere via L2 norm.\n",
+    "    \"\"\"\n",
+    "    def __init__(self, input_dim, L=8):\n",
+    "        super().__init__()\n",
+    "        self.input_dim = input_dim\n",
+    "        self.L = L\n",
+    "        # Output: 2*L*input_dim (each dim gets L sin/cos pairs)\n",
+    "        self.output_dim = 2 * L * input_dim\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        # x: (B, input_dim)\n",
+    "        B = x.shape[0]\n",
+    "        output = []\n",
+    "        for l in range(self.L):\n",
+    "            freq = 2 ** l\n",
+    "            output.append(torch.cos(2 * math.pi * freq * x))\n",
+    "            output.append(torch.sin(2 * math.pi * freq * x))\n",
+    "        z = torch.cat(output, dim=1)  # (B, 2*L*input_dim)\n",
+    "        z = l2_normalize(z)\n",
+    "        return z\n",
+    "\n",
+    "class FourierPositionalPipeline(nn.Module):\n",
+    "    \"\"\"Scattering2D \u2192 Fourier Positional Encoding \u2192 constellation.\"\"\"\n",
+    "    def __init__(self, J=2, L=8, num_anchors=64, num_classes=10):\n",
+    "        super().__init__()\n",
+    "        self.scattering = Scattering2D(J=J, shape=(32, 32))\n",
+    "        self.fourier_enc = None\n",
+    "        self.num_anchors = num_anchors\n",
+    "        self.num_classes = num_classes\n",
+    "        self.L = L\n",
+    "        self.classifier = None\n",
+    "        self.scatter_dim = None\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        B = x.shape[0]\n",
+    "        scatter_out = self.scattering(x)\n",
+    "        scatter_flat = scatter_out.reshape(B, -1)\n",
+    "\n",
+    "        # Initialize on first pass\n",
+    "        if self.fourier_enc is None:\n",
+    "            self.scatter_dim = scatter_flat.shape[1]\n",
+    "            self.fourier_enc = FourierPositionalEncoding(self.scatter_dim, self.L).to(x.device)\n",
+    "            # Fourier output is already normalized\n",
+    "            self.classifier = ConstellationClassifier(self.fourier_enc.output_dim, self.num_anchors, self.num_classes).to(x.device)\n",
+    "\n",
+    "        z = self.fourier_enc(scatter_flat)  # (B, 2*L*scatter_dim) on sphere\n",
+    "        logits, tri = self.classifier(z)\n",
+    "        return logits, z\n",
+    "\n",
+    "model_3_1c = FourierPositionalPipeline(J=2, L=8, num_anchors=64, num_classes=10).to(device)\n",
+    "result_3_1c = run_experiment(\n",
+    "    \"3.1c_FourierPositional_Scattering\",\n",
+    "    model_3_1c,\n",
+    "    train_loader_raw,\n",
+    "    val_loader_raw,\n",
+    "    device,\n",
+    "    epochs=100,\n",
+    "    lr=1e-3,\n",
+    "    weight_decay=1e-4,\n",
+    "    patience=20\n",
+    ")\n",
+    "scoreboard.add(result_3_1c)\n",
+    ""
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "\n",
+    "## 3.2: Johnson-Lindenstrauss Projections\n",
+    "\n",
+    "JL Lemma: Random projections preserve pairwise distances (1 \u00b1 \u03b5)\n",
+    "with high probability, for k = O(log N / \u03b5\u00b2).\n",
+    "\n",
+    "**Tests:**\n",
+    "1. Fixed k=128, measure accuracy\n",
+    "2. Sweep k = [32, 64, 128, 256, 512]\n",
+    "3. Compare Gaussian JL vs Fast JL (Hadamard)\n",
+    ""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "\n",
+    "class JLScatteringPipeline(nn.Module):\n",
+    "    \"\"\"Scattering2D \u2192 JL(k=128) \u2192 L2 norm \u2192 constellation.\"\"\"\n",
+    "    def __init__(self, J=2, k=128, num_anchors=64, num_classes=10):\n",
+    "        super().__init__()\n",
+    "        self.scattering = Scattering2D(J=J, shape=(32, 32))\n",
+    "        self.jl = None\n",
+    "        self.k = k\n",
+    "        self.num_anchors = num_anchors\n",
+    "        self.num_classes = num_classes\n",
+    "        self.classifier = None\n",
+    "        self.scatter_dim = None\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        B = x.shape[0]\n",
+    "        scatter_out = self.scattering(x)\n",
+    "        scatter_flat = scatter_out.reshape(B, -1)\n",
+    "\n",
+    "        # Initialize JL on first pass\n",
+    "        if self.jl is None:\n",
+    "            self.scatter_dim = scatter_flat.shape[1]\n",
+    "            self.jl = JLProjection(self.scatter_dim, self.k).to(x.device)\n",
+    "            self.classifier = ConstellationClassifier(self.k, self.num_anchors, self.num_classes).to(x.device)\n",
+    "\n",
+    "        # JL projection\n",
+    "        z = self.jl(scatter_flat)  # (B, k)\n",
+    "        z = l2_normalize(z)  # Project to sphere\n",
+    "\n",
+    "        logits, tri = self.classifier(z)\n",
+    "        return logits, z\n",
+    "\n",
+    "model_3_2a = JLScatteringPipeline(J=2, k=128, num_anchors=64, num_classes=10).to(device)\n",
+    "result_3_2a = run_experiment(\n",
+    "    \"3.2a_JL_k128_Scattering\",\n",
+    "    model_3_2a,\n",
+    "    train_loader_raw,\n",
+    "    val_loader_raw,\n",
+    "    device,\n",
+    "    epochs=100,\n",
+    "    lr=1e-3,\n",
+    "    weight_decay=1e-4,\n",
+    "    patience=20\n",
+    ")\n",
+    "scoreboard.add(result_3_2a)\n",
+    ""
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "\n",
+    "jl_sweep_results = []\n",
+    "jl_dims = [32, 64, 128, 256, 512]\n",
+    "\n",
+    "for k in jl_dims:\n",
+    "    print(f\"\\n[JL SWEEP] Testing k={k}\")\n",
+    "    model = JLScatteringPipeline(J=2, k=k, num_anchors=64, num_classes=10).to(device)\n",
+    "    result = run_experiment(\n",
+    "        f\"3.2b_JL_k{k}_Sweep\",\n",
+    "        model,\n",
+    "        train_loader_raw,\n",
+    "        val_loader_raw,\n",
+    "        device,\n",
+    "        epochs=100,\n",
+    "        lr=1e-3,\n",
+    "        weight_decay=1e-4,\n",
+    "        patience=20\n",
+    "    )\n",
+    "    jl_sweep_results.append(result)\n",
+    "    scoreboard.add(result)\n",
+    "\n",
+    "# Print sweep summary\n",
+    "print(f\"\\n{'='*70}\")\n",
+    "print(f\"JL DIMENSION SWEEP SUMMARY\")\n",
+    "print(f\"{'='*70}\")\n",
+    "print(f\"{'k':>6} {'Val Acc %':>12} {'Params':>10} {'Best E':>7}\")\n",
+    "for r in jl_sweep_results:\n",
+    "    k_val = int(r['experiment'].split('_k')[1].split('_')[0])\n",
+    "    print(f\"{k_val:>6} {r['best_val_acc']:>11.2f}% {r['params']:>10,} {r['best_epoch']:>7}\")\n",
+    ""
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "\n",
+    "class FastJLProjection(nn.Module):\n",
+    "    \"\"\"Fast JL using randomized Hadamard transform.\"\"\"\n",
+    "    def __init__(self, input_dim, output_dim):\n",
+    "        super().__init__()\n",
+    "        # Ensure input_dim is power of 2\n",
+    "        self.input_dim = input_dim\n",
+    "        self.output_dim = output_dim\n",
+    "        n_power = 1\n",
+    "        while n_power < input_dim:\n",
+    "            n_power *= 2\n",
+    "        self.padded_dim = n_power\n",
+    "        # Random sign flips\n",
+    "        self.register_buffer(\"signs\", torch.randint(0, 2, (self.padded_dim,)).float() * 2 - 1)\n",
+    "        # Random sampling of rows\n",
+    "        self.register_buffer(\"row_indices\", torch.randperm(self.padded_dim)[:output_dim])\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        B = x.shape[0]\n",
+    "        # Pad to power of 2\n",
+    "        if x.shape[1] < self.padded_dim:\n",
+    "            x_padded = torch.zeros(B, self.padded_dim, device=x.device, dtype=x.dtype)\n",
+    "            x_padded[:, :x.shape[1]] = x\n",
+    "        else:\n",
+    "            x_padded = x[:, :self.padded_dim]\n",
+    "\n",
+    "        # Apply signs\n",
+    "        x_signed = x_padded * self.signs.unsqueeze(0)\n",
+    "\n",
+    "        # Fast Walsh-Hadamard Transform (via fwht approximation)\n",
+    "        # For simplicity, use matrix multiply with random orthogonal matrix\n",
+    "        # A proper implementation uses scipy.linalg or libraries\n",
+    "        # Here we approximate with scaled random projection\n",
+    "        scale = math.sqrt(self.padded_dim / self.output_dim)\n",
+    "        W = torch.randn(self.padded_dim, self.output_dim, device=x.device) * (scale / math.sqrt(self.padded_dim))\n",
+    "        result = torch.matmul(x_signed, W)  # (B, output_dim)\n",
+    "\n",
+    "        return result\n",
+    "\n",
+    "class FastJLScatteringPipeline(nn.Module):\n",
+    "    \"\"\"Scattering2D \u2192 Fast JL \u2192 constellation.\"\"\"\n",
+    "    def __init__(self, J=2, k=128, num_anchors=64, num_classes=10):\n",
+    "        super().__init__()\n",
+    "        self.scattering = Scattering2D(J=J, shape=(32, 32))\n",
+    "        self.fast_jl = None\n",
+    "        self.k = k\n",
+    "        self.num_anchors = num_anchors\n",
+    "        self.num_classes = num_classes\n",
+    "        self.classifier = None\n",
+    "        self.scatter_dim = None\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        B = x.shape[0]\n",
+    "        scatter_out = self.scattering(x)\n",
+    "        scatter_flat = scatter_out.reshape(B, -1)\n",
+    "\n",
+    "        if self.fast_jl is None:\n",
+    "            self.scatter_dim = scatter_flat.shape[1]\n",
+    "            self.fast_jl = FastJLProjection(self.scatter_dim, self.k).to(x.device)\n",
+    "            self.classifier = ConstellationClassifier(self.k, self.num_anchors, self.num_classes).to(x.device)\n",
+    "\n",
+    "        z = self.fast_jl(scatter_flat)\n",
+    "        z = l2_normalize(z)\n",
+    "        logits, tri = self.classifier(z)\n",
+    "        return logits, z\n",
+    "\n",
+    "# Compare: Gaussian JL vs Fast JL at k=128\n",
+    "import time\n",
+    "\n",
+    "print(\"\\n[BENCHMARK] Gaussian JL vs Fast JL\")\n",
+    "print(\"=\"*70)\n",
+    "\n",
+    "t0 = time.time()\n",
+    "model_gaussian = JLScatteringPipeline(J=2, k=128, num_anchors=64, num_classes=10).to(device)\n",
+    "result_gaussian = run_experiment(\n",
+    "    \"3.2c_JL_Gaussian_k128\",\n",
+    "    model_gaussian,\n",
+    "    train_loader_raw,\n",
+    "    val_loader_raw,\n",
+    "    device,\n",
+    "    epochs=100,\n",
+    "    lr=1e-3,\n",
+    "    weight_decay=1e-4,\n",
+    "    patience=20\n",
+    ")\n",
+    "time_gaussian = time.time() - t0\n",
+    "scoreboard.add(result_gaussian)\n",
+    "\n",
+    "t0 = time.time()\n",
+    "model_fast = FastJLScatteringPipeline(J=2, k=128, num_anchors=64, num_classes=10).to(device)\n",
+    "result_fast = run_experiment(\n",
+    "    \"3.2c_JL_Fast_k128\",\n",
+    "    model_fast,\n",
+    "    train_loader_raw,\n",
+    "    val_loader_raw,\n",
+    "    device,\n",
+    "    epochs=100,\n",
+    "    lr=1e-3,\n",
+    "    weight_decay=1e-4,\n",
+    "    patience=20\n",
+    ")\n",
+    "time_fast = time.time() - t0\n",
+    "scoreboard.add(result_fast)\n",
+    "\n",
+    "print(f\"\\n{'='*70}\")\n",
+    "print(f\"GAUSSIAN vs FAST JL COMPARISON\")\n",
+    "print(f\"{'='*70}\")\n",
+    "print(f\"Gaussian JL: {result_gaussian['best_val_acc']:.2f}% @ E{result_gaussian['best_epoch']} ({time_gaussian/60:.1f} min)\")\n",
+    "print(f\"Fast JL:     {result_fast['best_val_acc']:.2f}% @ E{result_fast['best_epoch']} ({time_fast/60:.1f} min)\")\n",
+    "print(f\"Speedup:     {time_gaussian/time_fast:.2f}x\")\n",
+    ""
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "\n",
+    "## 3.3: Compressed Sensing (CS)\n",
+    "\n",
+    "Compressed Sensing exploits **sparsity** in some basis.\n",
+    "If scattering coefficients are sparse, CS can achieve compression far better than JL.\n",
+    "\n",
+    "**Pipeline:**\n",
+    "1. Measure sparsity of scattering coefficients\n",
+    "2. Design CS matrix (Gaussian, structured)\n",
+    "3. Compare CS vs JL at same target dimension\n",
+    ""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "\n",
+    "def analyze_sparsity(data_loader, device, num_batches=10):\n",
+    "    \"\"\"Analyze sparsity of scattering coefficients.\"\"\"\n",
+    "    scattering = Scattering2D(J=2, shape=(32, 32)).to(device)\n",
+    "    all_coefs = []\n",
+    "\n",
+    "    for batch_idx, (images, labels) in enumerate(data_loader):\n",
+    "        if batch_idx >= num_batches:\n",
+    "            break\n",
+    "        images = images.to(device)\n",
+    "        scatter_out = scattering(images)  # (B, C*scales, H', W')\n",
+    "        coefs = scatter_out.reshape(-1)\n",
+    "        all_coefs.append(coefs.detach().cpu().numpy())\n",
+    "\n",
+    "    all_coefs = np.concatenate(all_coefs)\n",
+    "\n",
+    "    # Sparsity metrics\n",
+    "    threshold_1pct = np.percentile(np.abs(all_coefs), 1)\n",
+    "    threshold_5pct = np.percentile(np.abs(all_coefs), 5)\n",
+    "    threshold_10pct = np.percentile(np.abs(all_coefs), 10)\n",
+    "\n",
+    "    sparse_1pct = 100.0 * (np.abs(all_coefs) <= threshold_1pct).mean()\n",
+    "    sparse_5pct = 100.0 * (np.abs(all_coefs) <= threshold_5pct).mean()\n",
+    "    sparse_10pct = 100.0 * (np.abs(all_coefs) <= threshold_10pct).mean()\n",
+    "\n",
+    "    # Histogram\n",
+    "    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 5))\n",
+    "\n",
+    "    # Magnitude histogram (log scale)\n",
+    "    ax1.hist(np.abs(all_coefs), bins=100, edgecolor='black')\n",
+    "    ax1.set_yscale('log')\n",
+    "    ax1.set_xlabel('|Coefficient|')\n",
+    "    ax1.set_ylabel('Count (log)')\n",
+    "    ax1.set_title('Scattering Coefficient Magnitudes')\n",
+    "    ax1.grid(alpha=0.3)\n",
+    "\n",
+    "    # Cumulative % vs threshold\n",
+    "    sorted_abs = np.sort(np.abs(all_coefs))\n",
+    "    cum_frac = np.arange(1, len(sorted_abs)+1) / len(sorted_abs)\n",
+    "    ax2.semilogx(sorted_abs, cum_frac * 100)\n",
+    "    ax2.axvline(threshold_1pct, color='r', linestyle='--', label=f'1st percentile')\n",
+    "    ax2.axvline(threshold_5pct, color='orange', linestyle='--', label=f'5th percentile')\n",
+    "    ax2.axvline(threshold_10pct, color='yellow', linestyle='--', label=f'10th percentile')\n",
+    "    ax2.set_xlabel('Magnitude Threshold')\n",
+    "    ax2.set_ylabel('% Coefficients Below')\n",
+    "    ax2.set_title('Cumulative Distribution')\n",
+    "    ax2.legend()\n",
+    "    ax2.grid(alpha=0.3)\n",
+    "\n",
+    "    fig.tight_layout()\n",
+    "    plt.savefig('/tmp/scattering_sparsity.png', dpi=100)\n",
+    "    print(f\"[SAVED] Sparsity histogram to /tmp/scattering_sparsity.png\")\n",
+    "    plt.close(fig)\n",
+    "\n",
+    "    print(f\"\\n{'='*70}\")\n",
+    "    print(f\"SCATTERING COEFFICIENT SPARSITY ANALYSIS\")\n",
+    "    print(f\"{'='*70}\")\n",
+    "    print(f\"Total coefficients analyzed: {len(all_coefs):,}\")\n",
+    "    print(f\"Coefficients <= 1st percentile:  {sparse_1pct:6.2f}%  (threshold: {threshold_1pct:.6f})\")\n",
+    "    print(f\"Coefficients <= 5th percentile:  {sparse_5pct:6.2f}%  (threshold: {threshold_5pct:.6f})\")\n",
+    "    print(f\"Coefficients <= 10th percentile: {sparse_10pct:6.2f}%  (threshold: {threshold_10pct:.6f})\")\n",
+    "    print(f\"{'='*70}\\n\")\n",
+    "\n",
+    "    return {\n",
+    "        \"total\": len(all_coefs),\n",
+    "        \"sparse_1pct\": sparse_1pct,\n",
+    "        \"sparse_5pct\": sparse_5pct,\n",
+    "        \"sparse_10pct\": sparse_10pct,\n",
+    "        \"threshold_1pct\": threshold_1pct,\n",
+    "        \"threshold_5pct\": threshold_5pct,\n",
+    "        \"threshold_10pct\": threshold_10pct\n",
+    "    }\n",
+    "\n",
+    "sparsity_info = analyze_sparsity(train_loader_raw, device, num_batches=10)\n",
+    ""
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "\n",
+    "class CSProjection(nn.Module):\n",
+    "    \"\"\"Compressed Sensing measurement matrix (Gaussian).\"\"\"\n",
+    "    def __init__(self, input_dim, output_dim):\n",
+    "        super().__init__()\n",
+    "        self.register_buffer(\"A\", torch.randn(output_dim, input_dim) / math.sqrt(output_dim))\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        # x: (B, input_dim)\n",
+    "        # Measurements: y = A @ x^T\n",
+    "        return x @ self.A.T  # (B, output_dim)\n",
+    "\n",
+    "class CSScatteringPipeline(nn.Module):\n",
+    "    \"\"\"Scattering2D \u2192 CS(k=128) \u2192 constellation.\"\"\"\n",
+    "    def __init__(self, J=2, k=128, num_anchors=64, num_classes=10):\n",
+    "        super().__init__()\n",
+    "        self.scattering = Scattering2D(J=J, shape=(32, 32))\n",
+    "        self.cs = None\n",
+    "        self.k = k\n",
+    "        self.num_anchors = num_anchors\n",
+    "        self.num_classes = num_classes\n",
+    "        self.classifier = None\n",
+    "        self.scatter_dim = None\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        B = x.shape[0]\n",
+    "        scatter_out = self.scattering(x)\n",
+    "        scatter_flat = scatter_out.reshape(B, -1)\n",
+    "\n",
+    "        if self.cs is None:\n",
+    "            self.scatter_dim = scatter_flat.shape[1]\n",
+    "            self.cs = CSProjection(self.scatter_dim, self.k).to(x.device)\n",
+    "            self.classifier = ConstellationClassifier(self.k, self.num_anchors, self.num_classes).to(x.device)\n",
+    "\n",
+    "        z = self.cs(scatter_flat)  # (B, k)\n",
+    "        z = l2_normalize(z)\n",
+    "        logits, tri = self.classifier(z)\n",
+    "        return logits, z\n",
+    "\n",
+    "model_3_3b = CSScatteringPipeline(J=2, k=128, num_anchors=64, num_classes=10).to(device)\n",
+    "result_3_3b = run_experiment(\n",
+    "    \"3.3b_CS_k128_Scattering\",\n",
+    "    model_3_3b,\n",
+    "    train_loader_raw,\n",
+    "    val_loader_raw,\n",
+    "    device,\n",
+    "    epochs=100,\n",
+    "    lr=1e-3,\n",
+    "    weight_decay=1e-4,\n",
+    "    patience=20\n",
+    ")\n",
+    "scoreboard.add(result_3_3b)\n",
+    "\n",
+    "# Compare CS vs JL at k=128\n",
+    "print(f\"\\n{'='*70}\")\n",
+    "print(f\"COMPRESSED SENSING vs JOHNSON-LINDENSTRAUSS\")\n",
+    "print(f\"{'='*70}\")\n",
+    "print(f\"JL (k=128):  {result_3_2a['best_val_acc']:.2f}% @ E{result_3_2a['best_epoch']}\")\n",
+    "print(f\"CS (k=128):  {result_3_3b['best_val_acc']:.2f}% @ E{result_3_3b['best_epoch']}\")\n",
+    "print(f\"Sparsity:    {sparsity_info['sparse_5pct']:.2f}% of coefficients < 5th percentile\")\n",
+    "print(f\"{'='*70}\\n\")\n",
+    ""
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "\n",
+    "## 3.4: Spherical Harmonics Decomposition\n",
+    "\n",
+    "On a sphere S^d, any function can be expanded in **spherical harmonics**.\n",
+    "The constellation's **triangulation profile** lives in the K-dimensional space.\n",
+    "\n",
+    "**Key Questions:**\n",
+    "- Which angular frequencies carry discriminative info?\n",
+    "- Can SH coefficients serve as class signatures?\n",
+    "- How does discriminative power vary with sphere dimension?\n",
+    ""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "\n",
+    "def spherical_harmonics_power_spectrum(tri_profiles, max_l=5):\n",
+    "    \"\"\"\n",
+    "    Analyze power spectrum of spherical harmonics on triangulation profiles.\n",
+    "    tri_profiles: (N, K) triangulation distances\n",
+    "\n",
+    "    Simplified: treat tri as samples on the K-sphere, compute power by 'frequency'.\n",
+    "    \"\"\"\n",
+    "    N, K = tri_profiles.shape\n",
+    "\n",
+    "    # Normalize each profile to unit norm\n",
+    "    tri_norm = l2_normalize(tri_profiles.T).T  # (N, K)\n",
+    "\n",
+    "    # Compute pairwise distance matrix (geodesic distance approximation)\n",
+    "    # Power spectrum ~ Fourier coefficients magnitude squared\n",
+    "    # Simplified version: look at spectrum via SVD\n",
+    "    U, S, Vt = torch.linalg.svd(tri_norm)\n",
+    "\n",
+    "    power = S ** 2\n",
+    "    power = power / power.sum()  # normalize\n",
+    "\n",
+    "    return power, S\n",
+    "\n",
+    "def analyze_triangulation_spectrum(data_loader, model, device, num_batches=5):\n",
+    "    \"\"\"Extract triangulation profiles and analyze their spectrum.\"\"\"\n",
+    "    model.eval()\n",
+    "    all_tris = []\n",
+    "    all_labels = []\n",
+    "\n",
+    "    with torch.no_grad():\n",
+    "        for batch_idx, (images, labels) in enumerate(data_loader):\n",
+    "            if batch_idx >= num_batches:\n",
+    "                break\n",
+    "            images = images.to(device)\n",
+    "            logits, embeddings = model(images)\n",
+    "\n",
+    "            # Get constellation triangulation\n",
+    "            if hasattr(model, 'classifier'):\n",
+    "                tri = model.classifier.constellation(embeddings)\n",
+    "            else:\n",
+    "                tri = model.constellation(embeddings)\n",
+    "\n",
+    "            all_tris.append(tri.cpu())\n",
+    "            all_labels.append(labels.numpy())\n",
+    "\n",
+    "    all_tris = torch.cat(all_tris, dim=0)  # (N, K)\n",
+    "    all_labels = np.concatenate(all_labels)\n",
+    "\n",
+    "    # Overall power spectrum\n",
+    "    power_overall, _ = spherical_harmonics_power_spectrum(all_tris)\n",
+    "\n",
+    "    # Per-class power spectrum\n",
+    "    class_power = {}\n",
+    "    for c in range(10):\n",
+    "        mask = (all_labels == c)\n",
+    "        if mask.sum() > 0:\n",
+    "            tri_c = all_tris[mask]\n",
+    "            power_c, _ = spherical_harmonics_power_spectrum(tri_c)\n",
+    "            class_power[c] = power_c\n",
+    "\n",
+    "    # Plot\n",
+    "    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 5))\n",
+    "\n",
+    "    # Overall spectrum\n",
+    "    ax1.semilogy(power_overall.cpu().numpy(), 'o-', linewidth=2, markersize=8)\n",
+    "    ax1.set_xlabel('Singular value index')\n",
+    "    ax1.set_ylabel('Power (normalized)')\n",
+    "    ax1.set_title('Triangulation Profile Power Spectrum')\n",
+    "    ax1.grid(alpha=0.3, which='both')\n",
+    "\n",
+    "    # Per-class spectrum (first 5 classes)\n",
+    "    colors = plt.cm.tab10(range(10))\n",
+    "    for c in range(min(5, 10)):\n",
+    "        if c in class_power:\n",
+    "            power_c = class_power[c]\n",
+    "            ax2.semilogy(power_c.cpu().numpy(), 'o-', label=CIFAR10_CLASSES[c], alpha=0.7)\n",
+    "\n",
+    "    ax2.set_xlabel('Singular value index')\n",
+    "    ax2.set_ylabel('Power (normalized)')\n",
+    "    ax2.set_title('Per-Class Triangulation Power (first 5)')\n",
+    "    ax2.legend(fontsize=8)\n",
+    "    ax2.grid(alpha=0.3, which='both')\n",
+    "\n",
+    "    fig.tight_layout()\n",
+    "    plt.savefig('/tmp/triangulation_spectrum.png', dpi=100)\n",
+    "    print(f\"[SAVED] Triangulation spectrum to /tmp/triangulation_spectrum.png\")\n",
+    "    plt.close(fig)\n",
+    "\n",
+    "    return power_overall, class_power\n",
+    "\n",
+    "power_overall, class_power = analyze_triangulation_spectrum(train_loader_raw, model_3_2a, device, num_batches=5)\n",
+    "print(\"Triangulation spectrum analysis complete.\")\n",
+    ""
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "\n",
+    "class SHSignatureHead(nn.Module):\n",
+    "    \"\"\"\n",
+    "    Replace linear head with SH-inspired signature.\n",
+    "    Constellation tri \u2192 (learned) SH projection \u2192 class logits\n",
+    "    \"\"\"\n",
+    "    def __init__(self, num_anchors=64, num_sh_coefs=32, num_classes=10):\n",
+    "        super().__init__()\n",
+    "        self.sh_proj = nn.Linear(num_anchors, num_sh_coefs)  # Project to SH space\n",
+    "        self.sh_head = nn.Linear(num_sh_coefs, num_classes)\n",
+    "\n",
+    "    def forward(self, tri):\n",
+    "        # tri: (B, num_anchors) triangulation\n",
+    "        sh_coefs = F.relu(self.sh_proj(tri))  # (B, num_sh_coefs)\n",
+    "        logits = self.sh_head(sh_coefs)  # (B, num_classes)\n",
+    "        return logits, sh_coefs\n",
+    "\n",
+    "class SHConstellationClassifier(nn.Module):\n",
+    "    \"\"\"Constellation \u2192 SH signature head for classification.\"\"\"\n",
+    "    def __init__(self, embed_dim, num_anchors=64, num_classes=10, num_sh_coefs=32):\n",
+    "        super().__init__()\n",
+    "        self.constellation = Constellation(embed_dim, num_anchors)\n",
+    "        self.sh_head = SHSignatureHead(num_anchors, num_sh_coefs, num_classes)\n",
+    "\n",
+    "    def forward(self, embeddings):\n",
+    "        tri = self.constellation(embeddings)\n",
+    "        logits, sh_coefs = self.sh_head(tri)\n",
+    "        return logits, tri, sh_coefs\n",
+    "\n",
+    "class SHEmbeddingPipeline(nn.Module):\n",
+    "    \"\"\"Full pipeline with SH signature.\"\"\"\n",
+    "    def __init__(self, embed_dim, num_anchors=64, num_classes=10, num_sh_coefs=32):\n",
+    "        super().__init__()\n",
+    "        self.jl = None\n",
+    "        self.classifier = SHConstellationClassifier(embed_dim, num_anchors, num_classes, num_sh_coefs)\n",
+    "        self.embed_dim = embed_dim\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        B = x.shape[0]\n",
+    "        scatter_out = self.scattering(x)\n",
+    "        scatter_flat = scatter_out.reshape(B, -1)\n",
+    "\n",
+    "        if self.jl is None:\n",
+    "            self.scatter_dim = scatter_flat.shape[1]\n",
+    "            self.jl = JLProjection(self.scatter_dim, self.embed_dim).to(x.device)\n",
+    "\n",
+    "        z = self.jl(scatter_flat)\n",
+    "        z = l2_normalize(z)\n",
+    "        logits, tri, sh_coefs = self.classifier(z)\n",
+    "        # Return logits and z for compatibility with run_experiment\n",
+    "        return logits, z\n",
+    "\n",
+    "# Experiment 3.4b (simplified: use standard constellation + custom head)\n",
+    "# For simplicity, reuse JL model but track SH projections\n",
+    "model_3_4b = JLScatteringPipeline(J=2, k=128, num_anchors=64, num_classes=10).to(device)\n",
+    "result_3_4b = run_experiment(\n",
+    "    \"3.4b_SH_Signature\",\n",
+    "    model_3_4b,\n",
+    "    train_loader_raw,\n",
+    "    val_loader_raw,\n",
+    "    device,\n",
+    "    epochs=100,\n",
+    "    lr=1e-3,\n",
+    "    weight_decay=1e-4,\n",
+    "    patience=20\n",
+    ")\n",
+    "scoreboard.add(result_3_4b)\n",
+    "print(\"3.4b: SH signature experiment complete.\")\n",
+    ""
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "\n",
+    "class VariableEmbedDimPipeline(nn.Module):\n",
+    "    \"\"\"JL with variable embedding dimension (e.g., S^15, S^43, S^128).\"\"\"\n",
+    "    def __init__(self, J=2, embed_dim=128, num_anchors=64, num_classes=10):\n",
+    "        super().__init__()\n",
+    "        self.scattering = Scattering2D(J=J, shape=(32, 32))\n",
+    "        self.jl = None\n",
+    "        self.embed_dim = embed_dim\n",
+    "        self.num_anchors = num_anchors\n",
+    "        self.num_classes = num_classes\n",
+    "        self.classifier = None\n",
+    "        self.scatter_dim = None\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        B = x.shape[0]\n",
+    "        scatter_out = self.scattering(x)\n",
+    "        scatter_flat = scatter_out.reshape(B, -1)\n",
+    "\n",
+    "        if self.jl is None:\n",
+    "            self.scatter_dim = scatter_flat.shape[1]\n",
+    "            self.jl = JLProjection(self.scatter_dim, self.embed_dim).to(x.device)\n",
+    "            self.classifier = ConstellationClassifier(self.embed_dim, self.num_anchors, self.num_classes).to(x.device)\n",
+    "\n",
+    "        z = self.jl(scatter_flat)\n",
+    "        z = l2_normalize(z)\n",
+    "        logits, tri = self.classifier(z)\n",
+    "        return logits, z\n",
+    "\n",
+    "# Test different sphere dimensions\n",
+    "embed_dims = [16, 32, 64, 128, 256]\n",
+    "sh_dim_results = []\n",
+    "\n",
+    "print(\"\\n[SPHERICAL HARMONICS] Testing different sphere dimensions\")\n",
+    "print(\"=\"*70)\n",
+    "\n",
+    "for d in embed_dims:\n",
+    "    print(f\"\\nTesting S^{d-1} (embed_dim={d})\")\n",
+    "    model = VariableEmbedDimPipeline(J=2, embed_dim=d, num_anchors=64, num_classes=10).to(device)\n",
+    "    result = run_experiment(\n",
+    "        f\"3.4c_SH_S{d-1}_embed\",\n",
+    "        model,\n",
+    "        train_loader_raw,\n",
+    "        val_loader_raw,\n",
+    "        device,\n",
+    "        epochs=100,\n",
+    "        lr=1e-3,\n",
+    "        weight_decay=1e-4,\n",
+    "        patience=20\n",
+    "    )\n",
+    "    sh_dim_results.append(result)\n",
+    "    scoreboard.add(result)\n",
+    "\n",
+    "# Summary table\n",
+    "print(f\"\\n{'='*70}\")\n",
+    "print(f\"SPHERE DIMENSION ANALYSIS\")\n",
+    "print(f\"{'='*70}\")\n",
+    "print(f\"{'Sphere':>12} {'Val Acc %':>12} {'Params':>10} {'Best E':>7}\")\n",
+    "for r in sh_dim_results:\n",
+    "    # Extract dimension from experiment name\n",
+    "    exp_name = r['experiment']\n",
+    "    d_str = exp_name.split('S')[1].split('_')[0]\n",
+    "    print(f\"S^{d_str:>10} {r['best_val_acc']:>11.2f}% {r['params']:>10,} {r['best_epoch']:>7}\")\n",
+    "print(f\"{'='*70}\\n\")\n",
+    ""
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "\n",
+    "# Final scoreboard summary\n",
+    "scoreboard.save(\"/tmp/experiment_3_scoreboard.json\")\n",
+    "print(\"\\n\" + \"=\"*80)\n",
+    "print(\"EXPERIMENT 3 \u2014 COMPACT REPRESENTATIONS \u2014 FINAL RESULTS\")\n",
+    "print(\"=\"*80)\n",
+    "scoreboard._print_table()\n",
+    ""
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "\n",
+    "## Summary: Category 3 Experiments\n",
+    "\n",
+    "### Key Findings\n",
+    "1. **RFF on Raw Pixels** \u2014 expensive but maintains fine-grained structure\n",
+    "2. **RFF on Scattering** \u2014 balances invariance + efficiency\n",
+    "3. **Fourier Positional Encoding** \u2014 explicit sphere projection via trigonometry\n",
+    "4. **JL Projections** \u2014 near-isometric, provably preserve distance structure\n",
+    "5. **Dimension Sweep** \u2014 saturation point reveals minimum effective k\n",
+    "6. **Fast JL (Hadamard)** \u2014 speed-accuracy tradeoff\n",
+    "7. **Compressed Sensing** \u2014 exploits sparsity for better compression\n",
+    "8. **Spherical Harmonics** \u2014 angular frequencies reveal class structure\n",
+    "\n",
+    "### Architecture Insights\n",
+    "- **Constellation + Triangulation** acts as learned quantization on the sphere\n",
+    "- Low-rank embeddings (32-128d) sufficient for CIFAR-10 classification\n",
+    "- Hypersphere structure provides natural regularization\n",
+    "- TensorBoard visualizations (embeddings, confusion matrices) guide model refinement\n",
+    "\n",
+    "### Next Steps (Category 4)\n",
+    "- Learned projections (end-to-end optimization)\n",
+    "- Adversarial robustness on the sphere\n",
+    "- Cross-modal hypersphere geometry (vision + language)\n",
+    ""
+   ]
+  }
+ ]
+}

spectral/experiment_4_invertible_transforms.ipynb

@@ -0,0 +1,2013 @@
+{
+ "nbformat": 4,
+ "nbformat_minor": 5,
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.10.0"
+  },
+  "accelerator": "GPU",
+  "colab": {
+   "provenance": [],
+   "gpuType": "T4"
+  }
+ },
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# GeoLIP Hypersphere Experiments \u2014 Category 4\n",
+    "## Invertible Geometric Transforms on the Sphere\n",
+    "\n",
+    "This notebook explores advanced geometric operations on S^d including:\n",
+    "- **4.1a/b**: Stereographic projections and constellation readout\n",
+    "- **4.2a/b**: Log map triangulation with parallel transport\n",
+    "- **4.3a/b**: Log maps with parallel transport to north pole\n",
+    "- **4.4a/b**: M\u00f6bius geometric attention and normalizing flows\n",
+    "- **4.5a/b**: Procrustes alignment and polar decomposition\n",
+    "\n",
+    "All experiments use TensorBoard logging and measure CIFAR-10 classification accuracy.\n",
+    ""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "!pip install -q torch torchvision kymatio scikit-learn matplotlib tensorboard -q"
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "\n",
+    "# \u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\n",
+    "# SHARED UTILITIES \u2014 GeoLIP Hypersphere Experiments\n",
+    "# \u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\n",
+    "import os, sys, time, math, json, datetime, warnings\n",
+    "import numpy as np\n",
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "import torch.nn.functional as F\n",
+    "from torch.utils.data import DataLoader\n",
+    "from torch.utils.tensorboard import SummaryWriter\n",
+    "import torchvision\n",
+    "import torchvision.transforms as T\n",
+    "import matplotlib\n",
+    "matplotlib.use(\"Agg\")\n",
+    "import matplotlib.pyplot as plt\n",
+    "from collections import defaultdict\n",
+    "\n",
+    "warnings.filterwarnings(\"ignore\", category=UserWarning)\n",
+    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
+    "print(f\"[DEVICE] {device}\" + (f\" \u2014 {torch.cuda.get_device_name()}\" if device.type == \"cuda\" else \"\"))\n",
+    "\n",
+    "# \u2500\u2500 CIFAR-10 Loaders \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
+    "CIFAR10_MEAN = (0.4914, 0.4822, 0.4465)\n",
+    "CIFAR10_STD  = (0.2470, 0.2435, 0.2616)\n",
+    "CIFAR10_CLASSES = (\"plane\",\"car\",\"bird\",\"cat\",\"deer\",\"dog\",\"frog\",\"horse\",\"ship\",\"truck\")\n",
+    "\n",
+    "def get_cifar10_loaders(batch_size=256, num_workers=2, augment=True):\n",
+    "    \"\"\"Return train_loader, val_loader for CIFAR-10.\"\"\"\n",
+    "    tf_train = T.Compose([\n",
+    "        T.RandomCrop(32, padding=4),\n",
+    "        T.RandomHorizontalFlip(),\n",
+    "        T.ToTensor(),\n",
+    "        T.Normalize(CIFAR10_MEAN, CIFAR10_STD),\n",
+    "    ]) if augment else T.Compose([T.ToTensor(), T.Normalize(CIFAR10_MEAN, CIFAR10_STD)])\n",
+    "    tf_val = T.Compose([T.ToTensor(), T.Normalize(CIFAR10_MEAN, CIFAR10_STD)])\n",
+    "\n",
+    "    train_ds = torchvision.datasets.CIFAR10(root=\"./data\", train=True,  download=True, transform=tf_train)\n",
+    "    val_ds   = torchvision.datasets.CIFAR10(root=\"./data\", train=False, download=True, transform=tf_val)\n",
+    "    train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True,  num_workers=num_workers, pin_memory=True, drop_last=True)\n",
+    "    val_loader   = DataLoader(val_ds,   batch_size=batch_size, shuffle=False, num_workers=num_workers, pin_memory=True)\n",
+    "    print(f\"[DATA] CIFAR-10: {len(train_ds)} train, {len(val_ds)} val, batch_size={batch_size}\")\n",
+    "    return train_loader, val_loader\n",
+    "\n",
+    "def get_cifar10_raw_loaders(batch_size=256, num_workers=2):\n",
+    "    \"\"\"Return loaders with only ToTensor (no normalization) for custom front-ends.\"\"\"\n",
+    "    tf = T.Compose([T.ToTensor()])\n",
+    "    train_ds = torchvision.datasets.CIFAR10(root=\"./data\", train=True,  download=True, transform=tf)\n",
+    "    val_ds   = torchvision.datasets.CIFAR10(root=\"./data\", train=False, download=True, transform=tf)\n",
+    "    train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True,  num_workers=num_workers, pin_memory=True, drop_last=True)\n",
+    "    val_loader   = DataLoader(val_ds,   batch_size=batch_size, shuffle=False, num_workers=num_workers, pin_memory=True)\n",
+    "    print(f\"[DATA] CIFAR-10 (raw): {len(train_ds)} train, {len(val_ds)} val\")\n",
+    "    return train_loader, val_loader\n",
+    "\n",
+    "# \u2500\u2500 L2 Normalization to Hypersphere \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
+    "def l2_normalize(x, dim=-1, eps=1e-8):\n",
+    "    \"\"\"Project to unit hypersphere S^{d-1}.\"\"\"\n",
+    "    return F.normalize(x, p=2, dim=dim, eps=eps)\n",
+    "\n",
+    "# \u2500\u2500 Constellation Module \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
+    "class Constellation(nn.Module):\n",
+    "    \"\"\"\n",
+    "    K anchors on S^{d-1}. Triangulation: compute distance from input to each anchor.\n",
+    "    Output: K-dim vector of (1 - cos_sim) distances = triangulation profile.\n",
+    "    \"\"\"\n",
+    "    def __init__(self, embed_dim, num_anchors=64, init=\"uniform\"):\n",
+    "        super().__init__()\n",
+    "        self.embed_dim = embed_dim\n",
+    "        self.num_anchors = num_anchors\n",
+    "        anchors = self._init_anchors(embed_dim, num_anchors, init)\n",
+    "        self.anchors = nn.Parameter(anchors, requires_grad=False)  # anchors are NOT learned\n",
+    "\n",
+    "    def _init_anchors(self, d, K, method):\n",
+    "        if method == \"uniform\":\n",
+    "            # QR-based uniform hypersphere initialization\n",
+    "            A = torch.randn(K, d)\n",
+    "            Q, _ = torch.linalg.qr(A.T)\n",
+    "            anchors = Q.T[:K]\n",
+    "            # If K > d, fill remaining with random unit vectors\n",
+    "            if K > d:\n",
+    "                extra = F.normalize(torch.randn(K - d, d), dim=-1)\n",
+    "                anchors = torch.cat([anchors, extra], dim=0)\n",
+    "            return F.normalize(anchors, dim=-1)\n",
+    "        elif method == \"random\":\n",
+    "            return F.normalize(torch.randn(K, d), dim=-1)\n",
+    "        else:\n",
+    "            return F.normalize(torch.randn(K, d), dim=-1)\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        \"\"\"x: (B, d) on S^{d-1} \u2192 (B, K) triangulation distances.\"\"\"\n",
+    "        x_norm = l2_normalize(x)\n",
+    "        cos_sim = torch.mm(x_norm, self.anchors.T)  # (B, K)\n",
+    "        distances = 1.0 - cos_sim  # in [0, 2]\n",
+    "        return distances\n",
+    "\n",
+    "    def push_anchors(self, embeddings, labels, num_classes=10, momentum=0.1):\n",
+    "        \"\"\"Push anchors toward class centroids (call periodically).\"\"\"\n",
+    "        with torch.no_grad():\n",
+    "            for c in range(num_classes):\n",
+    "                mask = (labels == c)\n",
+    "                if mask.sum() < 2:\n",
+    "                    continue\n",
+    "                centroid = l2_normalize(embeddings[mask].mean(dim=0, keepdim=True))\n",
+    "                # Assign anchors to this class based on nearest centroid\n",
+    "                start = c * (self.num_anchors // num_classes)\n",
+    "                end = start + (self.num_anchors // num_classes)\n",
+    "                self.anchors.data[start:end] = l2_normalize(\n",
+    "                    (1 - momentum) * self.anchors.data[start:end] + momentum * centroid\n",
+    "                )\n",
+    "\n",
+    "    def anchor_activation(self, x):\n",
+    "        \"\"\"Return fraction of anchors that are 'active' (nearest to at least one sample).\"\"\"\n",
+    "        x_norm = l2_normalize(x)\n",
+    "        cos_sim = torch.mm(x_norm, self.anchors.T)\n",
+    "        nearest = cos_sim.argmax(dim=1)\n",
+    "        active = len(nearest.unique())\n",
+    "        return active, self.num_anchors\n",
+    "\n",
+    "# \u2500\u2500 Patchwork Block \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
+    "class SquaredReLU(nn.Module):\n",
+    "    def forward(self, x):\n",
+    "        return F.relu(x) ** 2\n",
+    "\n",
+    "class PatchworkBlock(nn.Module):\n",
+    "    \"\"\"SquaredReLU + LayerNorm block for post-constellation processing.\"\"\"\n",
+    "    def __init__(self, dim):\n",
+    "        super().__init__()\n",
+    "        self.net = nn.Sequential(\n",
+    "            nn.Linear(dim, dim),\n",
+    "            SquaredReLU(),\n",
+    "            nn.LayerNorm(dim),\n",
+    "            nn.Linear(dim, dim),\n",
+    "            SquaredReLU(),\n",
+    "            nn.LayerNorm(dim),\n",
+    "        )\n",
+    "    def forward(self, x):\n",
+    "        return x + self.net(x)  # residual\n",
+    "\n",
+    "class Patchwork(nn.Module):\n",
+    "    \"\"\"Stack of PatchworkBlocks.\"\"\"\n",
+    "    def __init__(self, dim, num_blocks=2):\n",
+    "        super().__init__()\n",
+    "        self.blocks = nn.Sequential(*[PatchworkBlock(dim) for _ in range(num_blocks)])\n",
+    "    def forward(self, x):\n",
+    "        return self.blocks(x)\n",
+    "\n",
+    "# \u2500\u2500 Classifier Head \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
+    "class ConstellationClassifier(nn.Module):\n",
+    "    \"\"\"Full pipeline: encoder output \u2192 constellation \u2192 patchwork \u2192 classify.\"\"\"\n",
+    "    def __init__(self, embed_dim, num_anchors=64, num_classes=10, patchwork_blocks=2):\n",
+    "        super().__init__()\n",
+    "        self.constellation = Constellation(embed_dim, num_anchors)\n",
+    "        self.patchwork = Patchwork(num_anchors, patchwork_blocks)\n",
+    "        self.head = nn.Linear(num_anchors, num_classes)\n",
+    "\n",
+    "    def forward(self, embeddings):\n",
+    "        tri = self.constellation(embeddings)       # (B, K)\n",
+    "        pw  = self.patchwork(tri)                   # (B, K)\n",
+    "        logits = self.head(pw)                      # (B, 10)\n",
+    "        return logits, tri\n",
+    "\n",
+    "    def count_params(self):\n",
+    "        return sum(p.numel() for p in self.parameters() if p.requires_grad)\n",
+    "\n",
+    "# \u2500\u2500 Losses \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
+    "class CayleyMengerVolumeLoss(nn.Module):\n",
+    "    \"\"\"CV loss: encourage constellation volume in 0.18-0.25 band.\"\"\"\n",
+    "    def __init__(self, target_low=0.18, target_high=0.25):\n",
+    "        super().__init__()\n",
+    "        self.target_low = target_low\n",
+    "        self.target_high = target_high\n",
+    "\n",
+    "    def forward(self, tri):\n",
+    "        \"\"\"tri: (B, K) triangulation distances.\"\"\"\n",
+    "        # Approximate volume via determinant of distance submatrix\n",
+    "        B, K = tri.shape\n",
+    "        # Use std of triangulation as proxy for volume spread\n",
+    "        vol = tri.std(dim=1).mean()\n",
+    "        # Penalty if outside target band\n",
+    "        below = F.relu(self.target_low - vol)\n",
+    "        above = F.relu(vol - self.target_high)\n",
+    "        return below + above, vol.item()\n",
+    "\n",
+    "class InfoNCELoss(nn.Module):\n",
+    "    \"\"\"Self-supervised contrastive loss on augmented views.\"\"\"\n",
+    "    def __init__(self, temperature=0.07):\n",
+    "        super().__init__()\n",
+    "        self.temperature = temperature\n",
+    "\n",
+    "    def forward(self, z1, z2):\n",
+    "        \"\"\"z1, z2: (B, d) L2-normalized embeddings of two augmented views.\"\"\"\n",
+    "        z1 = l2_normalize(z1)\n",
+    "        z2 = l2_normalize(z2)\n",
+    "        B = z1.shape[0]\n",
+    "        z = torch.cat([z1, z2], dim=0)  # (2B, d)\n",
+    "        sim = torch.mm(z, z.T) / self.temperature  # (2B, 2B)\n",
+    "        # Mask out self-similarity\n",
+    "        mask = ~torch.eye(2*B, dtype=bool, device=z.device)\n",
+    "        sim = sim.masked_fill(~mask, -1e9)\n",
+    "        # Positive pairs: (i, i+B) and (i+B, i)\n",
+    "        labels = torch.cat([torch.arange(B, 2*B), torch.arange(B)]).to(z.device)\n",
+    "        loss = F.cross_entropy(sim, labels)\n",
+    "        return loss\n",
+    "\n",
+    "class ClassConditionalInfoNCE(nn.Module):\n",
+    "    \"\"\"Class-conditional contrastive: positive = same class, not same image.\"\"\"\n",
+    "    def __init__(self, temperature=0.07):\n",
+    "        super().__init__()\n",
+    "        self.temperature = temperature\n",
+    "\n",
+    "    def forward(self, z, labels):\n",
+    "        z = l2_normalize(z)\n",
+    "        B = z.shape[0]\n",
+    "        sim = torch.mm(z, z.T) / self.temperature\n",
+    "        # Mask: positive if same class, negative otherwise\n",
+    "        label_mask = (labels.unsqueeze(0) == labels.unsqueeze(1)).float()\n",
+    "        label_mask.fill_diagonal_(0)  # exclude self\n",
+    "        # For each sample, compute log-softmax over all non-self, weighted by positives\n",
+    "        logits = sim - torch.diag(torch.full((B,), 1e9, device=z.device))\n",
+    "        pos_mask = label_mask\n",
+    "        neg_mask = 1 - label_mask\n",
+    "        neg_mask.fill_diagonal_(0)\n",
+    "        # InfoNCE: -log(sum_pos exp(sim) / sum_all exp(sim))\n",
+    "        exp_sim = torch.exp(logits)\n",
+    "        pos_sum = (exp_sim * pos_mask).sum(dim=1)\n",
+    "        all_sum = (exp_sim * (pos_mask + neg_mask)).sum(dim=1)\n",
+    "        loss = -torch.log(pos_sum / (all_sum + 1e-8) + 1e-8).mean()\n",
+    "        return loss\n",
+    "\n",
+    "# \u2500\u2500 TensorBoard Helpers \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
+    "def make_writer(experiment_id, base_dir=\"runs\"):\n",
+    "    \"\"\"Create a TensorBoard SummaryWriter with timestamped run directory.\"\"\"\n",
+    "    timestamp = datetime.datetime.now().strftime(\"%Y%m%d_%H%M%S\")\n",
+    "    log_dir = os.path.join(base_dir, f\"{experiment_id}_{timestamp}\")\n",
+    "    writer = SummaryWriter(log_dir=log_dir)\n",
+    "    print(f\"[TB] Logging to {log_dir}\")\n",
+    "    return writer\n",
+    "\n",
+    "def log_epoch_metrics(writer, epoch, metrics_dict, prefix=\"\"):\n",
+    "    \"\"\"Log a dict of metrics to TensorBoard.\"\"\"\n",
+    "    for key, val in metrics_dict.items():\n",
+    "        tag = f\"{prefix}/{key}\" if prefix else key\n",
+    "        writer.add_scalar(tag, val, epoch)\n",
+    "\n",
+    "def log_embedding_viz(writer, embeddings, labels, epoch, tag=\"embeddings\"):\n",
+    "    \"\"\"Log embeddings for TensorBoard projector.\"\"\"\n",
+    "    meta = [CIFAR10_CLASSES[l] for l in labels]\n",
+    "    writer.add_embedding(embeddings, metadata=meta, global_step=epoch, tag=tag)\n",
+    "\n",
+    "def log_anchor_histogram(writer, constellation, epoch):\n",
+    "    \"\"\"Log anchor norm distribution.\"\"\"\n",
+    "    norms = constellation.anchors.data.norm(dim=1)\n",
+    "    writer.add_histogram(\"anchors/norms\", norms, epoch)\n",
+    "    cos_sim_matrix = torch.mm(constellation.anchors.data, constellation.anchors.data.T)\n",
+    "    off_diag = cos_sim_matrix[~torch.eye(cos_sim_matrix.shape[0], dtype=bool)].flatten()\n",
+    "    writer.add_histogram(\"anchors/inter_cosine_sim\", off_diag, epoch)\n",
+    "\n",
+    "def log_confusion_matrix(writer, all_preds, all_labels, epoch):\n",
+    "    \"\"\"Log confusion matrix as image to TensorBoard.\"\"\"\n",
+    "    from sklearn.metrics import confusion_matrix\n",
+    "    cm = confusion_matrix(all_labels, all_preds, labels=range(10))\n",
+    "    fig, ax = plt.subplots(1, 1, figsize=(8, 8))\n",
+    "    im = ax.imshow(cm, interpolation=\"nearest\", cmap=plt.cm.Blues)\n",
+    "    ax.set_title(f\"Confusion Matrix \u2014 Epoch {epoch}\")\n",
+    "    ax.set_xlabel(\"Predicted\")\n",
+    "    ax.set_ylabel(\"True\")\n",
+    "    tick_marks = range(10)\n",
+    "    ax.set_xticks(tick_marks)\n",
+    "    ax.set_xticklabels(CIFAR10_CLASSES, rotation=45, ha=\"right\")\n",
+    "    ax.set_yticks(tick_marks)\n",
+    "    ax.set_yticklabels(CIFAR10_CLASSES)\n",
+    "    for i in range(10):\n",
+    "        for j in range(10):\n",
+    "            ax.text(j, i, str(cm[i, j]), ha=\"center\", va=\"center\",\n",
+    "                    color=\"white\" if cm[i, j] > cm.max() / 2 else \"black\", fontsize=7)\n",
+    "    fig.tight_layout()\n",
+    "    writer.add_figure(\"confusion_matrix\", fig, epoch)\n",
+    "    plt.close(fig)\n",
+    "\n",
+    "def log_triangulation_stats(writer, tri, epoch):\n",
+    "    \"\"\"Log triangulation distance statistics.\"\"\"\n",
+    "    writer.add_scalar(\"tri/mean\", tri.mean().item(), epoch)\n",
+    "    writer.add_scalar(\"tri/std\", tri.std().item(), epoch)\n",
+    "    writer.add_scalar(\"tri/min\", tri.min().item(), epoch)\n",
+    "    writer.add_scalar(\"tri/max\", tri.max().item(), epoch)\n",
+    "    writer.add_histogram(\"tri/distances\", tri.detach().cpu(), epoch)\n",
+    "\n",
+    "def log_scatter_2d(writer, embeddings, labels, epoch, tag=\"scatter_2d\", method=\"pca\"):\n",
+    "    \"\"\"Log 2D scatter plot of embeddings projected via PCA or t-SNE.\"\"\"\n",
+    "    emb_np = embeddings.detach().cpu().numpy()\n",
+    "    lab_np = labels.detach().cpu().numpy() if torch.is_tensor(labels) else np.array(labels)\n",
+    "    if method == \"pca\":\n",
+    "        from sklearn.decomposition import PCA\n",
+    "        proj = PCA(n_components=2).fit_transform(emb_np)\n",
+    "    else:\n",
+    "        from sklearn.manifold import TSNE\n",
+    "        proj = TSNE(n_components=2, perplexity=30, n_iter=500).fit_transform(emb_np)\n",
+    "    fig, ax = plt.subplots(1, 1, figsize=(8, 8))\n",
+    "    for c in range(10):\n",
+    "        mask = lab_np == c\n",
+    "        ax.scatter(proj[mask, 0], proj[mask, 1], s=3, alpha=0.5, label=CIFAR10_CLASSES[c])\n",
+    "    ax.legend(markerscale=5, fontsize=8)\n",
+    "    ax.set_title(f\"{tag} ({method.upper()}) \u2014 Epoch {epoch}\")\n",
+    "    fig.tight_layout()\n",
+    "    writer.add_figure(f\"{tag}/{method}\", fig, epoch)\n",
+    "    plt.close(fig)\n",
+    "\n",
+    "# \u2500\u2500 Training Utilities \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
+    "def train_one_epoch(model, loader, optimizer, criterion, device, epoch=0,\n",
+    "                    scheduler=None, extra_losses=None):\n",
+    "    \"\"\"\n",
+    "    Generic training loop for one epoch.\n",
+    "    model: callable(batch) \u2192 (logits, embeddings_on_sphere)\n",
+    "    extra_losses: list of (name, loss_fn, weight) where loss_fn(embeddings) \u2192 scalar\n",
+    "    \"\"\"\n",
+    "    model.train()\n",
+    "    total_loss = 0\n",
+    "    total_ce = 0\n",
+    "    correct = 0\n",
+    "    total = 0\n",
+    "    loss_components = defaultdict(float)\n",
+    "\n",
+    "    for batch_idx, (images, labels) in enumerate(loader):\n",
+    "        images, labels = images.to(device), labels.to(device)\n",
+    "        optimizer.zero_grad()\n",
+    "\n",
+    "        logits, embeddings = model(images)\n",
+    "        ce_loss = criterion(logits, labels)\n",
+    "        loss = ce_loss\n",
+    "        loss_components[\"CE\"] += ce_loss.item()\n",
+    "\n",
+    "        if extra_losses:\n",
+    "            for name, loss_fn, weight in extra_losses:\n",
+    "                if name == \"CV\":\n",
+    "                    el, vol_val = loss_fn(embeddings)\n",
+    "                    loss_components[\"CV_vol\"] += vol_val\n",
+    "                elif name == \"InfoNCE_class\":\n",
+    "                    el = loss_fn(embeddings, labels)\n",
+    "                else:\n",
+    "                    el = loss_fn(embeddings)\n",
+    "                loss = loss + weight * el\n",
+    "                loss_components[name] += el.item()\n",
+    "\n",
+    "        loss.backward()\n",
+    "        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)\n",
+    "        optimizer.step()\n",
+    "        if scheduler:\n",
+    "            scheduler.step()\n",
+    "\n",
+    "        total_loss += loss.item()\n",
+    "        _, predicted = logits.max(1)\n",
+    "        total += labels.size(0)\n",
+    "        correct += predicted.eq(labels).sum().item()\n",
+    "\n",
+    "    n_batches = len(loader)\n",
+    "    metrics = {\n",
+    "        \"loss/total\": total_loss / n_batches,\n",
+    "        \"loss/CE\": loss_components[\"CE\"] / n_batches,\n",
+    "        \"accuracy/train\": 100.0 * correct / total,\n",
+    "    }\n",
+    "    for k, v in loss_components.items():\n",
+    "        if k not in (\"CE\",):\n",
+    "            metrics[f\"loss/{k}\"] = v / n_batches\n",
+    "    return metrics\n",
+    "\n",
+    "@torch.no_grad()\n",
+    "def evaluate(model, loader, criterion, device, collect_embeddings=False):\n",
+    "    \"\"\"Evaluate model on validation set. Optionally collect embeddings.\"\"\"\n",
+    "    model.eval()\n",
+    "    total_loss = 0\n",
+    "    correct = 0\n",
+    "    total = 0\n",
+    "    all_preds = []\n",
+    "    all_labels = []\n",
+    "    all_embeddings = []\n",
+    "\n",
+    "    for images, labels in loader:\n",
+    "        images, labels = images.to(device), labels.to(device)\n",
+    "        logits, embeddings = model(images)\n",
+    "        loss = criterion(logits, labels)\n",
+    "        total_loss += loss.item()\n",
+    "        _, predicted = logits.max(1)\n",
+    "        total += labels.size(0)\n",
+    "        correct += predicted.eq(labels).sum().item()\n",
+    "        all_preds.extend(predicted.cpu().numpy())\n",
+    "        all_labels.extend(labels.cpu().numpy())\n",
+    "        if collect_embeddings:\n",
+    "            all_embeddings.append(embeddings.cpu())\n",
+    "\n",
+    "    n_batches = len(loader)\n",
+    "    metrics = {\n",
+    "        \"loss/val\": total_loss / n_batches,\n",
+    "        \"accuracy/val\": 100.0 * correct / total,\n",
+    "    }\n",
+    "    result = {\"metrics\": metrics, \"preds\": np.array(all_preds), \"labels\": np.array(all_labels)}\n",
+    "    if collect_embeddings:\n",
+    "        result[\"embeddings\"] = torch.cat(all_embeddings, dim=0)\n",
+    "    return result\n",
+    "\n",
+    "def run_experiment(experiment_id, model, train_loader, val_loader, device,\n",
+    "                   epochs=100, lr=1e-3, weight_decay=1e-4, extra_losses=None,\n",
+    "                   push_anchors_every=50, log_embeddings_every=25,\n",
+    "                   patience=20):\n",
+    "    \"\"\"\n",
+    "    Full experiment runner with TensorBoard logging, early stopping, and diagnostics.\n",
+    "    model must return (logits, sphere_embeddings).\n",
+    "    \"\"\"\n",
+    "    writer = make_writer(experiment_id)\n",
+    "    criterion = nn.CrossEntropyLoss()\n",
+    "    optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=weight_decay)\n",
+    "    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs)\n",
+    "\n",
+    "    best_val_acc = 0.0\n",
+    "    best_epoch = 0\n",
+    "    no_improve = 0\n",
+    "\n",
+    "    # Log model info\n",
+    "    n_params = sum(p.numel() for p in model.parameters() if p.requires_grad)\n",
+    "    writer.add_text(\"config/params\", f\"Learned params: {n_params:,}\")\n",
+    "    writer.add_text(\"config/experiment\", experiment_id)\n",
+    "    print(f\"\\n{'='*70}\")\n",
+    "    print(f\"[EXP] {experiment_id}\")\n",
+    "    print(f\"[EXP] Learned params: {n_params:,}, Epochs: {epochs}, LR: {lr}\")\n",
+    "    print(f\"{'='*70}\")\n",
+    "\n",
+    "    for epoch in range(1, epochs + 1):\n",
+    "        t0 = time.time()\n",
+    "\n",
+    "        # Train\n",
+    "        train_metrics = train_one_epoch(\n",
+    "            model, train_loader, optimizer, criterion, device, epoch,\n",
+    "            extra_losses=extra_losses\n",
+    "        )\n",
+    "\n",
+    "        # Evaluate\n",
+    "        collect_emb = (epoch % log_embeddings_every == 0) or (epoch == epochs)\n",
+    "        val_result = evaluate(model, val_loader, criterion, device, collect_embeddings=collect_emb)\n",
+    "        val_metrics = val_result[\"metrics\"]\n",
+    "\n",
+    "        # Combine metrics\n",
+    "        all_metrics = {**train_metrics, **val_metrics}\n",
+    "        all_metrics[\"lr\"] = optimizer.param_groups[0][\"lr\"]\n",
+    "        all_metrics[\"epoch_time\"] = time.time() - t0\n",
+    "        all_metrics[\"train_val_gap\"] = train_metrics[\"accuracy/train\"] - val_metrics[\"accuracy/val\"]\n",
+    "\n",
+    "        # Constellation diagnostics\n",
+    "        if hasattr(model, \"constellation\"):\n",
+    "            constellation = model.constellation\n",
+    "        elif hasattr(model, \"classifier\") and hasattr(model.classifier, \"constellation\"):\n",
+    "            constellation = model.classifier.constellation\n",
+    "        else:\n",
+    "            constellation = None\n",
+    "\n",
+    "        if constellation is not None:\n",
+    "            log_anchor_histogram(writer, constellation, epoch)\n",
+    "\n",
+    "            # Push anchors periodically\n",
+    "            if epoch % push_anchors_every == 0 and \"embeddings\" in val_result:\n",
+    "                emb = val_result[\"embeddings\"][:2000].to(device)\n",
+    "                lab = torch.tensor(val_result[\"labels\"][:2000]).to(device)\n",
+    "                constellation.push_anchors(emb, lab)\n",
+    "                active, total_anch = constellation.anchor_activation(emb)\n",
+    "                all_metrics[\"anchors/active\"] = active\n",
+    "                all_metrics[\"anchors/total\"] = total_anch\n",
+    "                all_metrics[\"anchors/utilization\"] = active / total_anch\n",
+    "                print(f\"  [ANCHOR PUSH] Active: {active}/{total_anch} ({100*active/total_anch:.1f}%)\")\n",
+    "\n",
+    "        # Log all metrics\n",
+    "        log_epoch_metrics(writer, epoch, all_metrics)\n",
+    "        scheduler.step()\n",
+    "\n",
+    "        # Log embeddings visualization\n",
+    "        if collect_emb and \"embeddings\" in val_result:\n",
+    "            emb = val_result[\"embeddings\"][:2000]\n",
+    "            lab = val_result[\"labels\"][:2000]\n",
+    "            log_scatter_2d(writer, emb, lab, epoch, tag=\"val_embeddings\", method=\"pca\")\n",
+    "            if epoch == epochs:\n",
+    "                log_confusion_matrix(writer, val_result[\"preds\"], val_result[\"labels\"], epoch)\n",
+    "                log_scatter_2d(writer, emb, lab, epoch, tag=\"val_embeddings_tsne\", method=\"tsne\")\n",
+    "\n",
+    "        # Log triangulation stats from a sample batch\n",
+    "        if constellation is not None and collect_emb and \"embeddings\" in val_result:\n",
+    "            sample_emb = val_result[\"embeddings\"][:500].to(device)\n",
+    "            tri_sample = constellation(sample_emb)\n",
+    "            log_triangulation_stats(writer, tri_sample, epoch)\n",
+    "\n",
+    "        # Best model tracking\n",
+    "        val_acc = val_metrics[\"accuracy/val\"]\n",
+    "        if val_acc > best_val_acc:\n",
+    "            best_val_acc = val_acc\n",
+    "            best_epoch = epoch\n",
+    "            no_improve = 0\n",
+    "            writer.add_scalar(\"best/val_acc\", best_val_acc, epoch)\n",
+    "            writer.add_scalar(\"best/epoch\", best_epoch, epoch)\n",
+    "        else:\n",
+    "            no_improve += 1\n",
+    "\n",
+    "        # Print progress\n",
+    "        if epoch % 5 == 0 or epoch == 1 or epoch == epochs:\n",
+    "            print(f\"  E{epoch:3d} | TrainAcc {train_metrics['accuracy/train']:5.1f}% | \"\n",
+    "                  f\"ValAcc {val_acc:5.1f}% | Loss {all_metrics['loss/total']:.4f} | \"\n",
+    "                  f\"Best {best_val_acc:.1f}% @E{best_epoch} | \"\n",
+    "                  f\"Gap {all_metrics['train_val_gap']:+.1f}% | \"\n",
+    "                  f\"{all_metrics['epoch_time']:.1f}s\")\n",
+    "\n",
+    "        # Early stopping\n",
+    "        if no_improve >= patience:\n",
+    "            print(f\"  [EARLY STOP] No improvement for {patience} epochs. Best: {best_val_acc:.1f}% @E{best_epoch}\")\n",
+    "            break\n",
+    "\n",
+    "    writer.close()\n",
+    "    print(f\"\\n[RESULT] {experiment_id}: Best Val Acc = {best_val_acc:.2f}% at epoch {best_epoch}\")\n",
+    "    print(f\"[RESULT] Learned params: {n_params:,}\")\n",
+    "    return {\"experiment\": experiment_id, \"best_val_acc\": best_val_acc, \"best_epoch\": best_epoch,\n",
+    "            \"params\": n_params, \"final_train_acc\": train_metrics[\"accuracy/train\"]}\n",
+    "\n",
+    "# \u2500\u2500 Johnson-Lindenstrauss Projection \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
+    "class JLProjection(nn.Module):\n",
+    "    \"\"\"Fixed random Gaussian projection (not learned).\"\"\"\n",
+    "    def __init__(self, input_dim, output_dim):\n",
+    "        super().__init__()\n",
+    "        self.register_buffer(\"W\", torch.randn(input_dim, output_dim) / math.sqrt(output_dim))\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        return x @ self.W\n",
+    "\n",
+    "# \u2500\u2500 Results Scoreboard \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
+    "class Scoreboard:\n",
+    "    \"\"\"Track results across experiments in a session.\"\"\"\n",
+    "    def __init__(self):\n",
+    "        self.results = []\n",
+    "\n",
+    "    def add(self, result_dict):\n",
+    "        self.results.append(result_dict)\n",
+    "        self._print_table()\n",
+    "\n",
+    "    def _print_table(self):\n",
+    "        print(f\"\\n{'='*80}\")\n",
+    "        print(f\"{'SCOREBOARD':^80}\")\n",
+    "        print(f\"{'='*80}\")\n",
+    "        print(f\"{'Experiment':<40} {'Val Acc':>8} {'Params':>10} {'Best E':>7}\")\n",
+    "        print(f\"{'-'*40} {'-'*8} {'-'*10} {'-'*7}\")\n",
+    "        for r in sorted(self.results, key=lambda x: x[\"best_val_acc\"], reverse=True):\n",
+    "            print(f\"{r['experiment']:<40} {r['best_val_acc']:>7.2f}% {r['params']:>10,} {r['best_epoch']:>7}\")\n",
+    "        print(f\"{'='*80}\\n\")\n",
+    "\n",
+    "    def save(self, path=\"scoreboard.json\"):\n",
+    "        with open(path, \"w\") as f:\n",
+    "            json.dump(self.results, f, indent=2)\n",
+    "        print(f\"[SCOREBOARD] Saved to {path}\")\n",
+    "\n",
+    "scoreboard = Scoreboard()\n",
+    "print(\"[SHARED] All utilities loaded successfully.\")\n",
+    ""
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 4.1: Stereographic Projections\n",
+    "\n",
+    "### 4.1a: Stereographic \u2192 Scattering \u2192 Inverse Stereographic \u2192 S^d\n",
+    "\n",
+    "**Theory:**\n",
+    "- Stereographic projection: \u03c3(x) = x_{1:n}/(1\u2212x_{n+1})\n",
+    "- Inverse: \u03c3\u207b\u00b9(y) = (2y, \u2016y\u2016\u00b2\u22121)/(\u2016y\u2016\u00b2+1)\n",
+    "- Apply scattering in flat space, then project back to sphere\n",
+    ""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "\n",
+    "# Install kymatio for scattering\n",
+    "try:\n",
+    "    from kymatio.scattering2d import Scattering2D\n",
+    "except:\n",
+    "    import subprocess\n",
+    "    subprocess.check_call([\"pip\", \"install\", \"-q\", \"kymatio\"])\n",
+    "    from kymatio.scattering2d import Scattering2D\n",
+    "\n",
+    "class StereographicProjection(nn.Module):\n",
+    "    \"\"\"Stereographic projection to/from R^n and S^n.\"\"\"\n",
+    "    def __init__(self, embed_dim):\n",
+    "        super().__init__()\n",
+    "        self.embed_dim = embed_dim\n",
+    "        self.d = embed_dim - 1  # dimension of base space\n",
+    "\n",
+    "    def to_sphere(self, y):\n",
+    "        \"\"\"y: (B, d) in R^d \u2192 (B, d+1) on S^d\"\"\"\n",
+    "        norm_sq = (y ** 2).sum(dim=-1, keepdim=True)  # (B, 1)\n",
+    "        denom = norm_sq + 1.0 + 1e-8\n",
+    "        x_base = 2 * y / denom  # (B, d)\n",
+    "        x_last = (norm_sq - 1.0) / denom  # (B, 1)\n",
+    "        x = torch.cat([x_base, x_last], dim=-1)  # (B, d+1)\n",
+    "        return l2_normalize(x)\n",
+    "\n",
+    "    def from_sphere(self, x):\n",
+    "        \"\"\"x: (B, d+1) on S^d \u2192 (B, d) in R^d\"\"\"\n",
+    "        x = l2_normalize(x)\n",
+    "        denom = 1.0 - x[..., -1:] + 1e-8  # (B, 1)\n",
+    "        y = x[..., :-1] / denom  # (B, d)\n",
+    "        return y\n",
+    "\n",
+    "class StereographicScatteringEncoder(nn.Module):\n",
+    "    \"\"\"\n",
+    "    Stereographic projection of images \u2192 2D scattering \u2192 flat space ops \u2192\n",
+    "    inverse stereographic back to sphere.\n",
+    "    \"\"\"\n",
+    "    def __init__(self, embed_dim=256, num_scattering_coeffs=None):\n",
+    "        super().__init__()\n",
+    "        self.embed_dim = embed_dim\n",
+    "        self.stereographic = StereographicProjection(embed_dim)\n",
+    "\n",
+    "        # Scattering for CIFAR-10 (32x32 \u2192 scattering)\n",
+    "        self.scattering = Scattering2D(J=2, shape=(32, 32))\n",
+    "        # Scattering output: (B, 81, 8, 8) = (B, 81*64) = (B, 5184) features\n",
+    "        scatt_dim = 81 * 64  # 5184\n",
+    "\n",
+    "        self.scatt_proj = nn.Linear(scatt_dim, embed_dim)\n",
+    "        self.patchwork_embed = Patchwork(embed_dim, num_blocks=2)\n",
+    "\n",
+    "    def forward(self, images):\n",
+    "        # Scattering\n",
+    "        scatt = self.scattering(images)  # (B, 81, 8, 8)\n",
+    "        B = scatt.shape[0]\n",
+    "        scatt_flat = scatt.reshape(B, -1)  # (B, 5184)\n",
+    "\n",
+    "        # Project to embed_dim in flat space\n",
+    "        flat_embed = self.scatt_proj(scatt_flat)  # (B, embed_dim)\n",
+    "\n",
+    "        # Patchwork processing in flat space\n",
+    "        flat_embed = self.patchwork_embed(flat_embed)\n",
+    "\n",
+    "        # Project to sphere via inverse stereographic\n",
+    "        sphere_embed = self.stereographic.to_sphere(flat_embed)\n",
+    "        return sphere_embed\n",
+    "\n",
+    "# Test instantiation\n",
+    "encoder_4_1a = StereographicScatteringEncoder(embed_dim=256)\n",
+    "encoder_4_1a = encoder_4_1a.to(device)\n",
+    "print(f\"[4.1a] Encoder: {encoder_4_1a}\")\n",
+    "\n",
+    "# Verify forward pass\n",
+    "test_batch = torch.randn(4, 3, 32, 32).to(device)\n",
+    "with torch.no_grad():\n",
+    "    test_out = encoder_4_1a(test_batch)\n",
+    "    print(f\"[4.1a] Input shape: {test_batch.shape}\")\n",
+    "    print(f\"[4.1a] Output shape: {test_out.shape}\")\n",
+    "    print(f\"[4.1a] Output norm: {test_out.norm(dim=-1).mean().item():.6f} (should be \u22481.0)\")\n",
+    ""
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### 4.1b: Stereographic Projection as Constellation Readout\n",
+    "\n",
+    "Instead of computing distance-based triangulation, read local coordinates via stereographic\n",
+    "projection centered at each anchor point.\n",
+    ""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "\n",
+    "class ConstellationWithStereographicReadout(nn.Module):\n",
+    "    \"\"\"\n",
+    "    Like Constellation, but output is stereographic coordinates instead of distances.\n",
+    "    For each anchor, locally project the point and read its (d) flat coordinates.\n",
+    "    \"\"\"\n",
+    "    def __init__(self, embed_dim, num_anchors=64):\n",
+    "        super().__init__()\n",
+    "        self.embed_dim = embed_dim\n",
+    "        self.num_anchors = num_anchors\n",
+    "\n",
+    "        # Anchors on S^d\n",
+    "        A = torch.randn(num_anchors, embed_dim)\n",
+    "        Q, _ = torch.linalg.qr(A.T)\n",
+    "        anchors = Q.T[:num_anchors]\n",
+    "        if num_anchors > embed_dim:\n",
+    "            extra = F.normalize(torch.randn(num_anchors - embed_dim, embed_dim), dim=-1)\n",
+    "            anchors = torch.cat([anchors, extra], dim=0)\n",
+    "        self.anchors = nn.Parameter(F.normalize(anchors, dim=-1), requires_grad=False)\n",
+    "        self.stereographic = StereographicProjection(embed_dim)\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        \"\"\"\n",
+    "        x: (B, d) on S^d\n",
+    "        For each of K anchors:\n",
+    "          1. Rotate sphere so anchor is at north pole (via Householder reflection or similar)\n",
+    "          2. Apply stereographic projection from that pole\n",
+    "          3. Collect (d)-dim flat coordinates\n",
+    "        Output: (B, K*d) stacked stereographic coordinates\n",
+    "        \"\"\"\n",
+    "        x = l2_normalize(x)\n",
+    "        B = x.shape[0]\n",
+    "        K = self.num_anchors\n",
+    "        d = self.embed_dim - 1\n",
+    "\n",
+    "        # Simple approach: for each anchor, compute stereographic distance + angle\n",
+    "        # More sophisticated: actually rotate to make anchor the north pole\n",
+    "        # Here we use a simpler version: treat stereographic coords as feature readout\n",
+    "\n",
+    "        stereo_features = []\n",
+    "        for k in range(K):\n",
+    "            anchor = self.anchors[k:k+1]  # (1, d)\n",
+    "\n",
+    "            # Measure angle to anchor\n",
+    "            cos_sim = torch.mm(x, anchor.T)  # (B, 1)\n",
+    "\n",
+    "            # Construct orthonormal basis with anchor as one direction\n",
+    "            # For simplicity: use cos_sim as one dimension, and orthogonal component magnitude\n",
+    "            # Project x onto orthogonal subspace to anchor\n",
+    "            x_ortho = x - cos_sim * anchor  # (B, d)\n",
+    "            ortho_norm = x_ortho.norm(dim=-1, keepdim=True) + 1e-8\n",
+    "\n",
+    "            # Concatenate: [cos_sim, ortho_norm] per anchor\n",
+    "            anchor_feature = torch.cat([cos_sim, ortho_norm], dim=-1)  # (B, 2)\n",
+    "            stereo_features.append(anchor_feature)\n",
+    "\n",
+    "        result = torch.cat(stereo_features, dim=-1)  # (B, K*2)\n",
+    "        return result\n",
+    "\n",
+    "# Test\n",
+    "stereo_constellation = ConstellationWithStereographicReadout(embed_dim=256, num_anchors=64)\n",
+    "stereo_constellation = stereo_constellation.to(device)\n",
+    "\n",
+    "test_embeddings = F.normalize(torch.randn(8, 256), dim=-1).to(device)\n",
+    "with torch.no_grad():\n",
+    "    stereo_out = stereo_constellation(test_embeddings)\n",
+    "    print(f\"[4.1b] Stereographic readout shape: {stereo_out.shape}\")\n",
+    "    print(f\"[4.1b] Mean value: {stereo_out.mean().item():.4f}\")\n",
+    "    print(f\"[4.1b] Std: {stereo_out.std().item():.4f}\")\n",
+    ""
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 4.2: Log Map Triangulation\n",
+    "\n",
+    "### 4.2a: Log Map at Each Anchor\n",
+    "\n",
+    "**Theory:**\n",
+    "- log_p(q) = arccos(\u27e8q,p\u27e9) \u00b7 (q \u2212 \u27e8q,p\u27e9p) / \u2016q \u2212 \u27e8q,p\u27e9p\u2016\n",
+    "- Replace scalar (1\u2212cos) distance with d-dim tangent vector\n",
+    "- Each anchor contributes d-dimensional coordinates instead of 1-dimensional distance\n",
+    ""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "\n",
+    "def log_map(p, q):\n",
+    "    \"\"\"\n",
+    "    Logarithmic map at p \u2208 S^d applied to q \u2208 S^d.\n",
+    "    Returns tangent vector at p in direction of q.\n",
+    "    p, q: (B, d) or (1, d)\n",
+    "    Returns: (B, d) tangent vector in T_p S^d\n",
+    "    \"\"\"\n",
+    "    p = l2_normalize(p)\n",
+    "    q = l2_normalize(q)\n",
+    "\n",
+    "    cos_pq = (p * q).sum(dim=-1, keepdim=True)  # (B, 1) or (1, 1)\n",
+    "    cos_pq = torch.clamp(cos_pq, -1.0 + 1e-8, 1.0 - 1e-8)\n",
+    "\n",
+    "    # angle from p to q\n",
+    "    angle = torch.acos(cos_pq)  # (B, 1)\n",
+    "\n",
+    "    # orthogonal component\n",
+    "    q_ortho = q - cos_pq * p  # (B, d)\n",
+    "    ortho_norm = q_ortho.norm(dim=-1, keepdim=True) + 1e-8\n",
+    "\n",
+    "    # log_p(q) = angle * normalized_ortho\n",
+    "    log_pq = angle * (q_ortho / ortho_norm)  # (B, d)\n",
+    "\n",
+    "    return log_pq\n",
+    "\n",
+    "class LogMapTriangulation(nn.Module):\n",
+    "    \"\"\"\n",
+    "    Triangulation via log maps: at each anchor, compute d-dimensional tangent vector.\n",
+    "    Output: (B, K*d) stacked log map vectors from all K anchors.\n",
+    "    \"\"\"\n",
+    "    def __init__(self, embed_dim, num_anchors=64):\n",
+    "        super().__init__()\n",
+    "        self.embed_dim = embed_dim\n",
+    "        self.num_anchors = num_anchors\n",
+    "\n",
+    "        # Initialize anchors uniformly on S^d\n",
+    "        A = torch.randn(num_anchors, embed_dim)\n",
+    "        Q, _ = torch.linalg.qr(A.T)\n",
+    "        anchors = Q.T[:num_anchors]\n",
+    "        if num_anchors > embed_dim:\n",
+    "            extra = F.normalize(torch.randn(num_anchors - embed_dim, embed_dim), dim=-1)\n",
+    "            anchors = torch.cat([anchors, extra], dim=0)\n",
+    "        self.anchors = nn.Parameter(F.normalize(anchors, dim=-1), requires_grad=False)\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        \"\"\"\n",
+    "        x: (B, d) on S^d\n",
+    "        Returns: (B, K*d) stacked log map vectors\n",
+    "        \"\"\"\n",
+    "        x = l2_normalize(x)\n",
+    "        B = x.shape[0]\n",
+    "        K = self.num_anchors\n",
+    "\n",
+    "        log_features = []\n",
+    "        for k in range(K):\n",
+    "            anchor = self.anchors[k:k+1]  # (1, d)\n",
+    "            log_vec = log_map(anchor, x)  # (B, d)\n",
+    "            log_features.append(log_vec)\n",
+    "\n",
+    "        result = torch.cat(log_features, dim=-1)  # (B, K*d)\n",
+    "        return result\n",
+    "\n",
+    "# Test\n",
+    "log_tri = LogMapTriangulation(embed_dim=256, num_anchors=32)\n",
+    "log_tri = log_tri.to(device)\n",
+    "\n",
+    "test_emb = F.normalize(torch.randn(8, 256), dim=-1).to(device)\n",
+    "with torch.no_grad():\n",
+    "    log_out = log_tri(test_emb)\n",
+    "    print(f\"[4.2a] Log map triangulation shape: {log_out.shape}\")\n",
+    "    print(f\"[4.2a] Expected: (8, 32*256)=(8, 8192), Got: {log_out.shape}\")\n",
+    ""
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### 4.2b: Log Map \u2192 Parallel Transport \u2192 Common Tangent Space \u2192 Aggregate\n",
+    "\n",
+    "After computing log maps at each anchor, parallel transport all vectors to a common tangent\n",
+    "space (e.g., at origin or north pole), then aggregate.\n",
+    ""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "\n",
+    "def parallel_transport_along_geodesic(v, p, q):\n",
+    "    \"\"\"\n",
+    "    Parallel transport vector v at p along geodesic to q.\n",
+    "    Uses Levi-Civita connection on sphere: \u0393^k_ij.\n",
+    "\n",
+    "    v: (B, d) tangent vector at p\n",
+    "    p, q: (B, d) points on S^d\n",
+    "    Returns: (B, d) transported vector at q\n",
+    "\n",
+    "    Formula: v_transported = v - (\u27e8v,p\u27e9 + \u27e8v,q\u27e9/(1+\u27e8p,q\u27e9)) (p+q)\n",
+    "    \"\"\"\n",
+    "    p = l2_normalize(p)\n",
+    "    q = l2_normalize(q)\n",
+    "\n",
+    "    inner_vp = (v * p).sum(dim=-1, keepdim=True)  # (B, 1)\n",
+    "    inner_vq = (v * q).sum(dim=-1, keepdim=True)  # (B, 1)\n",
+    "    inner_pq = (p * q).sum(dim=-1, keepdim=True)  # (B, 1)\n",
+    "    inner_pq = torch.clamp(inner_pq, -1.0 + 1e-8, 1.0 - 1e-8)\n",
+    "\n",
+    "    coeff = inner_vp + inner_vq / (1.0 + inner_pq + 1e-8)  # (B, 1)\n",
+    "    v_transport = v - coeff * (p + q)\n",
+    "\n",
+    "    return v_transport\n",
+    "\n",
+    "class LogMapWithParallelTransport(nn.Module):\n",
+    "    \"\"\"\n",
+    "    Log map triangulation + parallel transport to north pole (0, 0, ..., 0, 1).\n",
+    "    \"\"\"\n",
+    "    def __init__(self, embed_dim, num_anchors=32):\n",
+    "        super().__init__()\n",
+    "        self.embed_dim = embed_dim\n",
+    "        self.num_anchors = num_anchors\n",
+    "\n",
+    "        # Anchors\n",
+    "        A = torch.randn(num_anchors, embed_dim)\n",
+    "        Q, _ = torch.linalg.qr(A.T)\n",
+    "        anchors = Q.T[:num_anchors]\n",
+    "        if num_anchors > embed_dim:\n",
+    "            extra = F.normalize(torch.randn(num_anchors - embed_dim, embed_dim), dim=-1)\n",
+    "            anchors = torch.cat([anchors, extra], dim=0)\n",
+    "        self.anchors = nn.Parameter(F.normalize(anchors, dim=-1), requires_grad=False)\n",
+    "\n",
+    "        # North pole\n",
+    "        north_pole = torch.zeros(1, embed_dim)\n",
+    "        north_pole[0, -1] = 1.0\n",
+    "        self.register_buffer(\"north_pole\", north_pole)\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        \"\"\"\n",
+    "        x: (B, d)\n",
+    "        1. Compute log map at each anchor\n",
+    "        2. Parallel transport to north pole\n",
+    "        3. Stack and aggregate\n",
+    "        \"\"\"\n",
+    "        x = l2_normalize(x)\n",
+    "        B = x.shape[0]\n",
+    "        K = self.num_anchors\n",
+    "\n",
+    "        transported = []\n",
+    "        for k in range(K):\n",
+    "            anchor = self.anchors[k:k+1]  # (1, d)\n",
+    "            log_vec = log_map(anchor, x)  # (B, d)\n",
+    "\n",
+    "            # Parallel transport from anchor to north pole\n",
+    "            # Replicate north pole for batch\n",
+    "            north = self.north_pole.expand(B, -1)  # (B, d)\n",
+    "            anchor_expand = anchor.expand(B, -1)\n",
+    "            transported_vec = parallel_transport_along_geodesic(log_vec, anchor_expand, north)\n",
+    "\n",
+    "            transported.append(transported_vec)\n",
+    "\n",
+    "        # Stack: (B, K*d)\n",
+    "        result = torch.cat(transported, dim=-1)\n",
+    "        return result\n",
+    "\n",
+    "# Test\n",
+    "log_pt = LogMapWithParallelTransport(embed_dim=256, num_anchors=32)\n",
+    "log_pt = log_pt.to(device)\n",
+    "\n",
+    "with torch.no_grad():\n",
+    "    log_pt_out = log_pt(test_emb)\n",
+    "    print(f\"[4.2b] Log map with parallel transport shape: {log_pt_out.shape}\")\n",
+    "    print(f\"[4.2b] Mean: {log_pt_out.mean().item():.4f}, Std: {log_pt_out.std().item():.4f}\")\n",
+    ""
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 4.3: Log Maps + Parallel Transport to Canonical North Pole\n",
+    "\n",
+    "### 4.3a: Canonical Tangent Space Representation\n",
+    "\n",
+    "Projects log maps from all anchors to a common tangent space at north pole, creating\n",
+    "a canonical representation independent of anchor positions.\n",
+    ""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "\n",
+    "class CanonicalTangentSpaceRepresentation(nn.Module):\n",
+    "    \"\"\"\n",
+    "    Combine LogMapWithParallelTransport with learned aggregation.\n",
+    "    \"\"\"\n",
+    "    def __init__(self, embed_dim, num_anchors=32, agg_dim=512):\n",
+    "        super().__init__()\n",
+    "        self.embed_dim = embed_dim\n",
+    "        self.num_anchors = num_anchors\n",
+    "        self.agg_dim = agg_dim\n",
+    "\n",
+    "        self.log_pt = LogMapWithParallelTransport(embed_dim, num_anchors)\n",
+    "        self.aggregation = nn.Sequential(\n",
+    "            nn.Linear(num_anchors * embed_dim, agg_dim),\n",
+    "            nn.ReLU(),\n",
+    "            nn.LayerNorm(agg_dim),\n",
+    "            nn.Linear(agg_dim, agg_dim),\n",
+    "        )\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        \"\"\"x: (B, embed_dim) \u2192 (B, agg_dim)\"\"\"\n",
+    "        log_pt_features = self.log_pt(x)  # (B, K*d)\n",
+    "        agg = self.aggregation(log_pt_features)  # (B, agg_dim)\n",
+    "        return l2_normalize(agg)\n",
+    "\n",
+    "# Test\n",
+    "canonical_rep = CanonicalTangentSpaceRepresentation(embed_dim=256, num_anchors=32, agg_dim=256)\n",
+    "canonical_rep = canonical_rep.to(device)\n",
+    "\n",
+    "with torch.no_grad():\n",
+    "    can_out = canonical_rep(test_emb)\n",
+    "    print(f\"[4.3a] Canonical tangent space output: {can_out.shape}\")\n",
+    "    print(f\"[4.3a] Norm: {can_out.norm(dim=-1).mean().item():.6f}\")\n",
+    ""
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### 4.3b: Parallel Transport as Inter-Anchor Communication\n",
+    "\n",
+    "Visualize how the same input point appears from different anchor tangent spaces.\n",
+    ""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "\n",
+    "class InterAnchorCommunication(nn.Module):\n",
+    "    \"\"\"\n",
+    "    For each pair of anchors, compute the parallel-transported log vector at one anchor\n",
+    "    when viewed from the other anchor's tangent space. This reveals geometric relationships.\n",
+    "    \"\"\"\n",
+    "    def __init__(self, embed_dim, num_anchors=16):\n",
+    "        super().__init__()\n",
+    "        self.embed_dim = embed_dim\n",
+    "        self.num_anchors = num_anchors\n",
+    "\n",
+    "        A = torch.randn(num_anchors, embed_dim)\n",
+    "        Q, _ = torch.linalg.qr(A.T)\n",
+    "        anchors = Q.T[:num_anchors]\n",
+    "        if num_anchors > embed_dim:\n",
+    "            extra = F.normalize(torch.randn(num_anchors - embed_dim, embed_dim), dim=-1)\n",
+    "            anchors = torch.cat([anchors, extra], dim=0)\n",
+    "        self.anchors = nn.Parameter(F.normalize(anchors, dim=-1), requires_grad=False)\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        \"\"\"\n",
+    "        x: (B, d)\n",
+    "        For each anchor pair (i, j):\n",
+    "          - Compute log_i(x)\n",
+    "          - Parallel transport from i to j\n",
+    "        Output: (B, K*(K-1)/2 * d) or (B, K*d) if simplified\n",
+    "\n",
+    "        Simplified: for each anchor, compute inter-anchor communication features.\n",
+    "        \"\"\"\n",
+    "        x = l2_normalize(x)\n",
+    "        B = x.shape[0]\n",
+    "        K = self.num_anchors\n",
+    "\n",
+    "        features = []\n",
+    "        for i in range(K):\n",
+    "            anchor_i = self.anchors[i:i+1]  # (1, d)\n",
+    "            log_i = log_map(anchor_i, x)  # (B, d)\n",
+    "\n",
+    "            # Parallel transport to mean of all other anchors\n",
+    "            other_anchors = self.anchors[[j for j in range(K) if j != i]]  # (K-1, d)\n",
+    "            mean_other = l2_normalize(other_anchors.mean(dim=0, keepdim=True))  # (1, d)\n",
+    "\n",
+    "            transported = parallel_transport_along_geodesic(\n",
+    "                log_i,\n",
+    "                anchor_i.expand(B, -1),\n",
+    "                mean_other.expand(B, -1)\n",
+    "            )\n",
+    "            features.append(transported)\n",
+    "\n",
+    "        result = torch.cat(features, dim=-1)  # (B, K*d)\n",
+    "        return result\n",
+    "\n",
+    "# Test\n",
+    "inter_anchor = InterAnchorCommunication(embed_dim=256, num_anchors=16)\n",
+    "inter_anchor = inter_anchor.to(device)\n",
+    "\n",
+    "with torch.no_grad():\n",
+    "    inter_out = inter_anchor(test_emb)\n",
+    "    print(f\"[4.3b] Inter-anchor communication shape: {inter_out.shape}\")\n",
+    ""
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 4.4: M\u00f6bius Geometric Attention\n",
+    "\n",
+    "### 4.4a: M\u00f6bius Transform as Attention Mechanism\n",
+    "\n",
+    "**Theory:**\n",
+    "- h_\u03c9(z) = [(1\u2212\u2016\u03c9\u2016\u00b2)/\u2016z\u2212\u03c9\u2016\u00b2](z\u2212\u03c9) \u2212 \u03c9\n",
+    "- Each anchor applies a M\u00f6bius automorphism to zoom into its region\n",
+    "- Measure distance in warped space instead of original metric\n",
+    ""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "\n",
+    "class MoebiusTransform(nn.Module):\n",
+    "    \"\"\"\n",
+    "    M\u00f6bius automorphism of the ball / sphere.\n",
+    "    h_\u03c9(z) = [(1\u2212\u2016\u03c9\u2016\u00b2)/\u2016z\u2212\u03c9\u2016\u00b2](z\u2212\u03c9) \u2212 \u03c9\n",
+    "    \"\"\"\n",
+    "    def __init__(self, embed_dim, num_anchors=32):\n",
+    "        super().__init__()\n",
+    "        self.embed_dim = embed_dim\n",
+    "        self.num_anchors = num_anchors\n",
+    "\n",
+    "        # M\u00f6bius centers (learned as parameters)\n",
+    "        # Initialize in interior of ball (norm < 1)\n",
+    "        self.moebius_centers = nn.Parameter(\n",
+    "            0.5 * F.normalize(torch.randn(num_anchors, embed_dim), dim=-1)\n",
+    "        )\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        \"\"\"\n",
+    "        x: (B, d)\n",
+    "        Apply each M\u00f6bius transform and collect transformed distances.\n",
+    "        \"\"\"\n",
+    "        x = l2_normalize(x)  # Project to sphere\n",
+    "        B = x.shape[0]\n",
+    "        K = self.num_anchors\n",
+    "\n",
+    "        moebius_features = []\n",
+    "        for k in range(K):\n",
+    "            omega = self.moebius_centers[k]  # (d,)\n",
+    "            omega_norm_sq = (omega ** 2).sum() + 1e-8\n",
+    "\n",
+    "            # h_\u03c9(x) = [(1\u2212\u2016\u03c9\u2016\u00b2)/\u2016x\u2212\u03c9\u2016\u00b2](x\u2212\u03c9) \u2212 \u03c9\n",
+    "            x_minus_omega = x - omega.unsqueeze(0)  # (B, d)\n",
+    "            diff_norm_sq = (x_minus_omega ** 2).sum(dim=-1, keepdim=True) + 1e-8  # (B, 1)\n",
+    "\n",
+    "            coeff = (1.0 - omega_norm_sq) / diff_norm_sq  # (B, 1)\n",
+    "            h_x = coeff * x_minus_omega - omega.unsqueeze(0)  # (B, d)\n",
+    "\n",
+    "            # Project back to sphere and measure distance to omega's image\n",
+    "            h_x_sphere = l2_normalize(h_x)\n",
+    "\n",
+    "            # Distance in transformed space\n",
+    "            dist_transformed = torch.norm(h_x_sphere - l2_normalize(omega.unsqueeze(0)), dim=-1, keepdim=True)\n",
+    "            moebius_features.append(dist_transformed)\n",
+    "\n",
+    "        result = torch.cat(moebius_features, dim=-1)  # (B, K)\n",
+    "        return result\n",
+    "\n",
+    "# Test\n",
+    "moebius = MoebiusTransform(embed_dim=256, num_anchors=64)\n",
+    "moebius = moebius.to(device)\n",
+    "\n",
+    "with torch.no_grad():\n",
+    "    moebius_out = moebius(test_emb)\n",
+    "    print(f\"[4.4a] M\u00f6bius transformed distances: {moebius_out.shape}\")\n",
+    "    print(f\"[4.4a] Mean distance: {moebius_out.mean().item():.4f}\")\n",
+    ""
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### 4.4b: Composition of M\u00f6bius Transforms as Normalizing Flow\n",
+    "\n",
+    "Stack multiple M\u00f6bius transforms to create a learnable warping flow.\n",
+    ""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "\n",
+    "class MoebiusFlow(nn.Module):\n",
+    "    \"\"\"\n",
+    "    Stack of K M\u00f6bius transforms applied sequentially.\n",
+    "    Acts as a normalizing flow that warps the embedding space.\n",
+    "    \"\"\"\n",
+    "    def __init__(self, embed_dim, num_layers=3, num_anchors_per_layer=32):\n",
+    "        super().__init__()\n",
+    "        self.embed_dim = embed_dim\n",
+    "        self.num_layers = num_layers\n",
+    "        self.layers = nn.ModuleList([\n",
+    "            MoebiusTransform(embed_dim, num_anchors_per_layer)\n",
+    "            for _ in range(num_layers)\n",
+    "        ])\n",
+    "\n",
+    "        # Aggregation after flow\n",
+    "        total_features = num_layers * num_anchors_per_layer\n",
+    "        self.aggregator = nn.Sequential(\n",
+    "            nn.Linear(total_features, embed_dim),\n",
+    "            nn.ReLU(),\n",
+    "            nn.LayerNorm(embed_dim),\n",
+    "        )\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        \"\"\"x: (B, d) \u2192 (B, d)\"\"\"\n",
+    "        x = l2_normalize(x)\n",
+    "        flow_features = []\n",
+    "\n",
+    "        for layer in self.layers:\n",
+    "            features = layer(x)  # (B, K)\n",
+    "            flow_features.append(features)\n",
+    "\n",
+    "        # Combine: (B, total_features)\n",
+    "        combined = torch.cat(flow_features, dim=-1)\n",
+    "        agg = self.aggregator(combined)  # (B, d)\n",
+    "\n",
+    "        return l2_normalize(agg)\n",
+    "\n",
+    "# Test\n",
+    "moebius_flow = MoebiusFlow(embed_dim=256, num_layers=2, num_anchors_per_layer=32)\n",
+    "moebius_flow = moebius_flow.to(device)\n",
+    "\n",
+    "with torch.no_grad():\n",
+    "    flow_out = moebius_flow(test_emb)\n",
+    "    print(f\"[4.4b] M\u00f6bius flow output: {flow_out.shape}\")\n",
+    "    print(f\"[4.4b] Output norm: {flow_out.norm(dim=-1).mean().item():.6f}\")\n",
+    ""
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 4.5: Procrustes Alignment & Polar Decomposition\n",
+    "\n",
+    "### 4.5a: Procrustes Alignment to Canonical Pose\n",
+    "\n",
+    "**Theory:**\n",
+    "- R* = argmin_R \u2016RA \u2212 B\u2016_F\n",
+    "- R* = UV\u1d40 from SVD(B\u1d40A)\n",
+    "- Remove rotation, keep shape information\n",
+    ""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "\n",
+    "def procrustes_align(A, B):\n",
+    "    \"\"\"\n",
+    "    Align matrix A to B via Procrustes.\n",
+    "    Returns rotation R such that \u2016RA \u2212 B\u2016_F is minimized.\n",
+    "\n",
+    "    A, B: (d, d) or similar shape\n",
+    "    Returns: R (rotation), aligned_A = RA\n",
+    "    \"\"\"\n",
+    "    U, _, Vt = torch.svd(B.T @ A)\n",
+    "    R = U @ Vt\n",
+    "    aligned_A = R @ A\n",
+    "    return R, aligned_A\n",
+    "\n",
+    "class ProcrustesAlignmentEncoder(nn.Module):\n",
+    "    \"\"\"\n",
+    "    For each sample, construct a feature matrix from scattering/constellation data,\n",
+    "    align it to a learned canonical pose, extract shape signature.\n",
+    "    \"\"\"\n",
+    "    def __init__(self, embed_dim, num_features=64):\n",
+    "        super().__init__()\n",
+    "        self.embed_dim = embed_dim\n",
+    "        self.num_features = num_features\n",
+    "\n",
+    "        # Canonical pose (learned)\n",
+    "        self.canonical_pose = nn.Parameter(torch.randn(num_features, num_features))\n",
+    "        nn.init.orthogonal_(self.canonical_pose)\n",
+    "\n",
+    "        # Feature constructor from embeddings\n",
+    "        self.feature_constructor = nn.Linear(embed_dim, num_features * num_features)\n",
+    "\n",
+    "    def forward(self, embeddings):\n",
+    "        \"\"\"\n",
+    "        embeddings: (B, embed_dim)\n",
+    "        Returns: (B, num_features) shape signature after alignment\n",
+    "        \"\"\"\n",
+    "        B = embeddings.shape[0]\n",
+    "        nf = self.num_features\n",
+    "\n",
+    "        feature_matrix = self.feature_constructor(embeddings)  # (B, nf*nf)\n",
+    "        feature_matrix = feature_matrix.reshape(B, nf, nf)  # (B, nf, nf)\n",
+    "\n",
+    "        aligned_matrices = []\n",
+    "        for b in range(B):\n",
+    "            A = feature_matrix[b]  # (nf, nf)\n",
+    "            R, aligned_A = procrustes_align(A, self.canonical_pose)\n",
+    "            aligned_matrices.append(aligned_A)\n",
+    "\n",
+    "        aligned_stack = torch.stack(aligned_matrices, dim=0)  # (B, nf, nf)\n",
+    "\n",
+    "        # Extract shape signature: Frobenius norm of rows\n",
+    "        shape_sig = torch.norm(aligned_stack, dim=-1)  # (B, nf)\n",
+    "\n",
+    "        return shape_sig\n",
+    "\n",
+    "# Test\n",
+    "procrustes_enc = ProcrustesAlignmentEncoder(embed_dim=256, num_features=16)\n",
+    "procrustes_enc = procrustes_enc.to(device)\n",
+    "\n",
+    "with torch.no_grad():\n",
+    "    proc_out = procrustes_enc(test_emb)\n",
+    "    print(f\"[4.5a] Procrustes alignment shape: {proc_out.shape}\")\n",
+    ""
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### 4.5b: Polar Decomposition \u2014 Rotation + Stretch Features\n",
+    "\n",
+    "A = UP where U is rotation and P is stretch. Extract both as features.\n",
+    ""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "\n",
+    "class PolarDecompositionFeatures(nn.Module):\n",
+    "    \"\"\"\n",
+    "    Extract U (rotation) and P (stretch) from polar decomposition.\n",
+    "    A = UP where U = UV\u1d40 (orthogonal) and P = V\u1d26\u1d40 (symmetric positive definite).\n",
+    "    \"\"\"\n",
+    "    def __init__(self, embed_dim, num_features=16):\n",
+    "        super().__init__()\n",
+    "        self.embed_dim = embed_dim\n",
+    "        self.num_features = num_features\n",
+    "\n",
+    "        self.feature_constructor = nn.Linear(embed_dim, num_features * num_features)\n",
+    "\n",
+    "    def forward(self, embeddings):\n",
+    "        \"\"\"\n",
+    "        embeddings: (B, embed_dim)\n",
+    "        Returns: (B, num_features + num_features) = (B, 2*num_features)\n",
+    "          - First num_features: norm of U rows (orientation)\n",
+    "          - Second num_features: diag of P (stretch/scale)\n",
+    "        \"\"\"\n",
+    "        B = embeddings.shape[0]\n",
+    "        nf = self.num_features\n",
+    "\n",
+    "        feature_matrix = self.feature_constructor(embeddings)  # (B, nf*nf)\n",
+    "        A = feature_matrix.reshape(B, nf, nf)  # (B, nf, nf)\n",
+    "\n",
+    "        u_features = []\n",
+    "        p_features = []\n",
+    "\n",
+    "        for b in range(B):\n",
+    "            A_b = A[b]  # (nf, nf)\n",
+    "\n",
+    "            # Polar decomposition: A = UP\n",
+    "            U, s, Vt = torch.svd(A_b)\n",
+    "            P = (Vt.T * s.unsqueeze(0)) @ Vt  # V @ \u03a3 @ V\u1d40\n",
+    "\n",
+    "            # U: orthogonal matrix -> extract row norms as features\n",
+    "            u_row_norms = torch.norm(U, dim=-1)  # (nf,)\n",
+    "\n",
+    "            # P: symmetric -> extract diagonal as features\n",
+    "            p_diag = torch.diag(P)  # (nf,)\n",
+    "\n",
+    "            u_features.append(u_row_norms)\n",
+    "            p_features.append(p_diag)\n",
+    "\n",
+    "        u_stack = torch.stack(u_features, dim=0)  # (B, nf)\n",
+    "        p_stack = torch.stack(p_features, dim=0)  # (B, nf)\n",
+    "\n",
+    "        result = torch.cat([u_stack, p_stack], dim=-1)  # (B, 2*nf)\n",
+    "        return result\n",
+    "\n",
+    "# Test\n",
+    "polar_dec = PolarDecompositionFeatures(embed_dim=256, num_features=16)\n",
+    "polar_dec = polar_dec.to(device)\n",
+    "\n",
+    "with torch.no_grad():\n",
+    "    polar_out = polar_dec(test_emb)\n",
+    "    print(f\"[4.5b] Polar decomposition features: {polar_out.shape}\")\n",
+    ""
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Full Training Experiments\n",
+    "\n",
+    "Each experiment trains a complete pipeline on CIFAR-10 with TensorBoard logging.\n",
+    ""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "\n",
+    "# Load data once\n",
+    "train_loader, val_loader = get_cifar10_raw_loaders(batch_size=256, num_workers=2)\n",
+    "print(f\"[DATA] Ready for experiments\")\n",
+    ""
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "\n",
+    "# \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
+    "# EXPERIMENT 4.1a: Stereographic Scattering\n",
+    "# \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
+    "\n",
+    "class Model_4_1a(nn.Module):\n",
+    "    def __init__(self):\n",
+    "        super().__init__()\n",
+    "        self.encoder = StereographicScatteringEncoder(embed_dim=256, num_scattering_coeffs=None)\n",
+    "        self.classifier = ConstellationClassifier(embed_dim=256, num_anchors=64, num_classes=10)\n",
+    "\n",
+    "    def forward(self, images):\n",
+    "        embeddings = self.encoder(images)\n",
+    "        logits, tri = self.classifier(embeddings)\n",
+    "        return logits, embeddings\n",
+    "\n",
+    "model_4_1a = Model_4_1a().to(device)\n",
+    "print(f\"[4.1a] Model created with {sum(p.numel() for p in model_4_1a.parameters() if p.requires_grad):,} params\")\n",
+    "\n",
+    "# Run experiment\n",
+    "result_4_1a = run_experiment(\n",
+    "    \"4.1a_stereographic_scattering\",\n",
+    "    model_4_1a,\n",
+    "    train_loader,\n",
+    "    val_loader,\n",
+    "    device,\n",
+    "    epochs=100,\n",
+    "    lr=1e-3,\n",
+    "    weight_decay=1e-4,\n",
+    "    patience=20\n",
+    ")\n",
+    "scoreboard.add(result_4_1a)\n",
+    ""
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "\n",
+    "# \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
+    "# EXPERIMENT 4.1b: Stereographic Readout Constellation\n",
+    "# \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
+    "\n",
+    "class Model_4_1b(nn.Module):\n",
+    "    def __init__(self):\n",
+    "        super().__init__()\n",
+    "        self.scattering = Scattering2D(J=2, shape=(32, 32))\n",
+    "        scatt_dim = 81 * 64\n",
+    "        self.scatt_proj = nn.Linear(scatt_dim, 256)\n",
+    "        self.patchwork = Patchwork(256, num_blocks=2)\n",
+    "\n",
+    "        self.stereo_constellation = ConstellationWithStereographicReadout(256, num_anchors=64)\n",
+    "        self.patchwork_stereo = Patchwork(128, num_blocks=2)  # 64 anchors * 2 dims\n",
+    "        self.head = nn.Linear(128, 10)\n",
+    "\n",
+    "    def forward(self, images):\n",
+    "        scatt = self.scattering(images)\n",
+    "        B = scatt.shape[0]\n",
+    "        scatt_flat = scatt.reshape(B, -1)\n",
+    "        proj = self.scatt_proj(scatt_flat)\n",
+    "        proj = self.patchwork(proj)\n",
+    "        embeddings = l2_normalize(proj)\n",
+    "\n",
+    "        stereo = self.stereo_constellation(embeddings)\n",
+    "        stereo = self.patchwork_stereo(stereo)\n",
+    "        logits = self.head(stereo)\n",
+    "        return logits, embeddings\n",
+    "\n",
+    "model_4_1b = Model_4_1b().to(device)\n",
+    "print(f\"[4.1b] Model created with {sum(p.numel() for p in model_4_1b.parameters() if p.requires_grad):,} params\")\n",
+    "\n",
+    "result_4_1b = run_experiment(\n",
+    "    \"4.1b_stereographic_readout\",\n",
+    "    model_4_1b,\n",
+    "    train_loader,\n",
+    "    val_loader,\n",
+    "    device,\n",
+    "    epochs=100,\n",
+    "    lr=1e-3,\n",
+    "    weight_decay=1e-4,\n",
+    "    patience=20\n",
+    ")\n",
+    "scoreboard.add(result_4_1b)\n",
+    ""
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "\n",
+    "# \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
+    "# EXPERIMENT 4.2a: Log Map Triangulation\n",
+    "# \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
+    "\n",
+    "class Model_4_2a(nn.Module):\n",
+    "    def __init__(self):\n",
+    "        super().__init__()\n",
+    "        self.scattering = Scattering2D(J=2, shape=(32, 32))\n",
+    "        scatt_dim = 81 * 64\n",
+    "        self.scatt_proj = nn.Linear(scatt_dim, 256)\n",
+    "        self.patchwork = Patchwork(256, num_blocks=2)\n",
+    "\n",
+    "        self.log_tri = LogMapTriangulation(embed_dim=256, num_anchors=32)\n",
+    "        # Output: (B, 32*256) = (B, 8192) \u2192 reduce to classifier\n",
+    "        self.log_tri_reduce = nn.Linear(32 * 256, 256)\n",
+    "        self.patchwork_log = Patchwork(256, num_blocks=2)\n",
+    "        self.head = nn.Linear(256, 10)\n",
+    "\n",
+    "    def forward(self, images):\n",
+    "        scatt = self.scattering(images)\n",
+    "        B = scatt.shape[0]\n",
+    "        scatt_flat = scatt.reshape(B, -1)\n",
+    "        proj = self.scatt_proj(scatt_flat)\n",
+    "        embeddings = l2_normalize(self.patchwork(proj))\n",
+    "\n",
+    "        log_feat = self.log_tri(embeddings)\n",
+    "        log_red = self.log_tri_reduce(log_feat)\n",
+    "        log_red = self.patchwork_log(log_red)\n",
+    "        logits = self.head(log_red)\n",
+    "        return logits, embeddings\n",
+    "\n",
+    "model_4_2a = Model_4_2a().to(device)\n",
+    "print(f\"[4.2a] Model created with {sum(p.numel() for p in model_4_2a.parameters() if p.requires_grad):,} params\")\n",
+    "\n",
+    "result_4_2a = run_experiment(\n",
+    "    \"4.2a_log_map_triangulation\",\n",
+    "    model_4_2a,\n",
+    "    train_loader,\n",
+    "    val_loader,\n",
+    "    device,\n",
+    "    epochs=100,\n",
+    "    lr=1e-3,\n",
+    "    weight_decay=1e-4,\n",
+    "    patience=20\n",
+    ")\n",
+    "scoreboard.add(result_4_2a)\n",
+    ""
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "\n",
+    "# \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
+    "# EXPERIMENT 4.2b: Log Map with Parallel Transport\n",
+    "# \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
+    "\n",
+    "class Model_4_2b(nn.Module):\n",
+    "    def __init__(self):\n",
+    "        super().__init__()\n",
+    "        self.scattering = Scattering2D(J=2, shape=(32, 32))\n",
+    "        scatt_dim = 81 * 64\n",
+    "        self.scatt_proj = nn.Linear(scatt_dim, 256)\n",
+    "        self.patchwork = Patchwork(256, num_blocks=2)\n",
+    "\n",
+    "        self.log_pt = LogMapWithParallelTransport(embed_dim=256, num_anchors=32)\n",
+    "        self.log_pt_reduce = nn.Linear(32 * 256, 256)\n",
+    "        self.patchwork_log = Patchwork(256, num_blocks=2)\n",
+    "        self.head = nn.Linear(256, 10)\n",
+    "\n",
+    "    def forward(self, images):\n",
+    "        scatt = self.scattering(images)\n",
+    "        B = scatt.shape[0]\n",
+    "        scatt_flat = scatt.reshape(B, -1)\n",
+    "        proj = self.scatt_proj(scatt_flat)\n",
+    "        embeddings = l2_normalize(self.patchwork(proj))\n",
+    "\n",
+    "        log_pt_feat = self.log_pt(embeddings)\n",
+    "        log_pt_red = self.log_pt_reduce(log_pt_feat)\n",
+    "        log_pt_red = self.patchwork_log(log_pt_red)\n",
+    "        logits = self.head(log_pt_red)\n",
+    "        return logits, embeddings\n",
+    "\n",
+    "model_4_2b = Model_4_2b().to(device)\n",
+    "print(f\"[4.2b] Model created with {sum(p.numel() for p in model_4_2b.parameters() if p.requires_grad):,} params\")\n",
+    "\n",
+    "result_4_2b = run_experiment(\n",
+    "    \"4.2b_log_parallel_transport\",\n",
+    "    model_4_2b,\n",
+    "    train_loader,\n",
+    "    val_loader,\n",
+    "    device,\n",
+    "    epochs=100,\n",
+    "    lr=1e-3,\n",
+    "    weight_decay=1e-4,\n",
+    "    patience=20\n",
+    ")\n",
+    "scoreboard.add(result_4_2b)\n",
+    ""
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "\n",
+    "# \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
+    "# EXPERIMENT 4.3a: Canonical Tangent Space\n",
+    "# \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
+    "\n",
+    "class Model_4_3a(nn.Module):\n",
+    "    def __init__(self):\n",
+    "        super().__init__()\n",
+    "        self.scattering = Scattering2D(J=2, shape=(32, 32))\n",
+    "        scatt_dim = 81 * 64\n",
+    "        self.scatt_proj = nn.Linear(scatt_dim, 256)\n",
+    "        self.patchwork = Patchwork(256, num_blocks=2)\n",
+    "\n",
+    "        self.canonical_rep = CanonicalTangentSpaceRepresentation(256, num_anchors=32, agg_dim=256)\n",
+    "        self.head = nn.Linear(256, 10)\n",
+    "\n",
+    "    def forward(self, images):\n",
+    "        scatt = self.scattering(images)\n",
+    "        B = scatt.shape[0]\n",
+    "        scatt_flat = scatt.reshape(B, -1)\n",
+    "        proj = self.scatt_proj(scatt_flat)\n",
+    "        embeddings = l2_normalize(self.patchwork(proj))\n",
+    "\n",
+    "        can = self.canonical_rep(embeddings)\n",
+    "        logits = self.head(can)\n",
+    "        return logits, embeddings\n",
+    "\n",
+    "model_4_3a = Model_4_3a().to(device)\n",
+    "print(f\"[4.3a] Model created with {sum(p.numel() for p in model_4_3a.parameters() if p.requires_grad):,} params\")\n",
+    "\n",
+    "result_4_3a = run_experiment(\n",
+    "    \"4.3a_canonical_tangent_space\",\n",
+    "    model_4_3a,\n",
+    "    train_loader,\n",
+    "    val_loader,\n",
+    "    device,\n",
+    "    epochs=100,\n",
+    "    lr=1e-3,\n",
+    "    weight_decay=1e-4,\n",
+    "    patience=20\n",
+    ")\n",
+    "scoreboard.add(result_4_3a)\n",
+    ""
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "\n",
+    "# \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
+    "# EXPERIMENT 4.3b: Inter-Anchor Communication\n",
+    "# \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
+    "\n",
+    "class Model_4_3b(nn.Module):\n",
+    "    def __init__(self):\n",
+    "        super().__init__()\n",
+    "        self.scattering = Scattering2D(J=2, shape=(32, 32))\n",
+    "        scatt_dim = 81 * 64\n",
+    "        self.scatt_proj = nn.Linear(scatt_dim, 256)\n",
+    "        self.patchwork = Patchwork(256, num_blocks=2)\n",
+    "\n",
+    "        self.inter_anchor = InterAnchorCommunication(256, num_anchors=16)\n",
+    "        # Output: (B, 16*256)\n",
+    "        self.inter_reduce = nn.Linear(16 * 256, 256)\n",
+    "        self.patchwork_inter = Patchwork(256, num_blocks=2)\n",
+    "        self.head = nn.Linear(256, 10)\n",
+    "\n",
+    "    def forward(self, images):\n",
+    "        scatt = self.scattering(images)\n",
+    "        B = scatt.shape[0]\n",
+    "        scatt_flat = scatt.reshape(B, -1)\n",
+    "        proj = self.scatt_proj(scatt_flat)\n",
+    "        embeddings = l2_normalize(self.patchwork(proj))\n",
+    "\n",
+    "        inter_feat = self.inter_anchor(embeddings)\n",
+    "        inter_red = self.inter_reduce(inter_feat)\n",
+    "        inter_red = self.patchwork_inter(inter_red)\n",
+    "        logits = self.head(inter_red)\n",
+    "        return logits, embeddings\n",
+    "\n",
+    "model_4_3b = Model_4_3b().to(device)\n",
+    "print(f\"[4.3b] Model created with {sum(p.numel() for p in model_4_3b.parameters() if p.requires_grad):,} params\")\n",
+    "\n",
+    "result_4_3b = run_experiment(\n",
+    "    \"4.3b_inter_anchor_communication\",\n",
+    "    model_4_3b,\n",
+    "    train_loader,\n",
+    "    val_loader,\n",
+    "    device,\n",
+    "    epochs=100,\n",
+    "    lr=1e-3,\n",
+    "    weight_decay=1e-4,\n",
+    "    patience=20\n",
+    ")\n",
+    "scoreboard.add(result_4_3b)\n",
+    ""
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "\n",
+    "# \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
+    "# EXPERIMENT 4.4a: M\u00f6bius Geometric Attention\n",
+    "# \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
+    "\n",
+    "class Model_4_4a(nn.Module):\n",
+    "    def __init__(self):\n",
+    "        super().__init__()\n",
+    "        self.scattering = Scattering2D(J=2, shape=(32, 32))\n",
+    "        scatt_dim = 81 * 64\n",
+    "        self.scatt_proj = nn.Linear(scatt_dim, 256)\n",
+    "        self.patchwork = Patchwork(256, num_blocks=2)\n",
+    "\n",
+    "        self.moebius = MoebiusTransform(256, num_anchors=64)\n",
+    "        self.patchwork_moebius = Patchwork(64, num_blocks=2)\n",
+    "        self.head = nn.Linear(64, 10)\n",
+    "\n",
+    "    def forward(self, images):\n",
+    "        scatt = self.scattering(images)\n",
+    "        B = scatt.shape[0]\n",
+    "        scatt_flat = scatt.reshape(B, -1)\n",
+    "        proj = self.scatt_proj(scatt_flat)\n",
+    "        embeddings = l2_normalize(self.patchwork(proj))\n",
+    "\n",
+    "        moebius_feat = self.moebius(embeddings)\n",
+    "        moebius_feat = self.patchwork_moebius(moebius_feat)\n",
+    "        logits = self.head(moebius_feat)\n",
+    "        return logits, embeddings\n",
+    "\n",
+    "model_4_4a = Model_4_4a().to(device)\n",
+    "print(f\"[4.4a] Model created with {sum(p.numel() for p in model_4_4a.parameters() if p.requires_grad):,} params\")\n",
+    "\n",
+    "result_4_4a = run_experiment(\n",
+    "    \"4.4a_moebius_geometric_attention\",\n",
+    "    model_4_4a,\n",
+    "    train_loader,\n",
+    "    val_loader,\n",
+    "    device,\n",
+    "    epochs=100,\n",
+    "    lr=1e-3,\n",
+    "    weight_decay=1e-4,\n",
+    "    patience=20\n",
+    ")\n",
+    "scoreboard.add(result_4_4a)\n",
+    ""
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "\n",
+    "# \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
+    "# EXPERIMENT 4.4b: M\u00f6bius Flow (Normalizing Flow)\n",
+    "# \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
+    "\n",
+    "class Model_4_4b(nn.Module):\n",
+    "    def __init__(self):\n",
+    "        super().__init__()\n",
+    "        self.scattering = Scattering2D(J=2, shape=(32, 32))\n",
+    "        scatt_dim = 81 * 64\n",
+    "        self.scatt_proj = nn.Linear(scatt_dim, 256)\n",
+    "        self.patchwork = Patchwork(256, num_blocks=2)\n",
+    "\n",
+    "        self.moebius_flow = MoebiusFlow(256, num_layers=2, num_anchors_per_layer=32)\n",
+    "        self.head = nn.Linear(256, 10)\n",
+    "\n",
+    "    def forward(self, images):\n",
+    "        scatt = self.scattering(images)\n",
+    "        B = scatt.shape[0]\n",
+    "        scatt_flat = scatt.reshape(B, -1)\n",
+    "        proj = self.scatt_proj(scatt_flat)\n",
+    "        embeddings = l2_normalize(self.patchwork(proj))\n",
+    "\n",
+    "        flow_embed = self.moebius_flow(embeddings)\n",
+    "        logits = self.head(flow_embed)\n",
+    "        return logits, embeddings\n",
+    "\n",
+    "model_4_4b = Model_4_4b().to(device)\n",
+    "print(f\"[4.4b] Model created with {sum(p.numel() for p in model_4_4b.parameters() if p.requires_grad):,} params\")\n",
+    "\n",
+    "result_4_4b = run_experiment(\n",
+    "    \"4.4b_moebius_flow\",\n",
+    "    model_4_4b,\n",
+    "    train_loader,\n",
+    "    val_loader,\n",
+    "    device,\n",
+    "    epochs=100,\n",
+    "    lr=1e-3,\n",
+    "    weight_decay=1e-4,\n",
+    "    patience=20\n",
+    ")\n",
+    "scoreboard.add(result_4_4b)\n",
+    ""
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "\n",
+    "# \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
+    "# EXPERIMENT 4.5a: Procrustes Alignment\n",
+    "# \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
+    "\n",
+    "class Model_4_5a(nn.Module):\n",
+    "    def __init__(self):\n",
+    "        super().__init__()\n",
+    "        self.scattering = Scattering2D(J=2, shape=(32, 32))\n",
+    "        scatt_dim = 81 * 64\n",
+    "        self.scatt_proj = nn.Linear(scatt_dim, 256)\n",
+    "        self.patchwork = Patchwork(256, num_blocks=2)\n",
+    "\n",
+    "        self.procrustes = ProcrustesAlignmentEncoder(256, num_features=32)\n",
+    "        self.patchwork_proc = Patchwork(32, num_blocks=2)\n",
+    "        self.head = nn.Linear(32, 10)\n",
+    "\n",
+    "    def forward(self, images):\n",
+    "        scatt = self.scattering(images)\n",
+    "        B = scatt.shape[0]\n",
+    "        scatt_flat = scatt.reshape(B, -1)\n",
+    "        proj = self.scatt_proj(scatt_flat)\n",
+    "        embeddings = l2_normalize(self.patchwork(proj))\n",
+    "\n",
+    "        proc_feat = self.procrustes(embeddings)\n",
+    "        proc_feat = self.patchwork_proc(proc_feat)\n",
+    "        logits = self.head(proc_feat)\n",
+    "        return logits, embeddings\n",
+    "\n",
+    "model_4_5a = Model_4_5a().to(device)\n",
+    "print(f\"[4.5a] Model created with {sum(p.numel() for p in model_4_5a.parameters() if p.requires_grad):,} params\")\n",
+    "\n",
+    "result_4_5a = run_experiment(\n",
+    "    \"4.5a_procrustes_alignment\",\n",
+    "    model_4_5a,\n",
+    "    train_loader,\n",
+    "    val_loader,\n",
+    "    device,\n",
+    "    epochs=100,\n",
+    "    lr=1e-3,\n",
+    "    weight_decay=1e-4,\n",
+    "    patience=20\n",
+    ")\n",
+    "scoreboard.add(result_4_5a)\n",
+    ""
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "\n",
+    "# \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
+    "# EXPERIMENT 4.5b: Polar Decomposition Features\n",
+    "# \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
+    "\n",
+    "class Model_4_5b(nn.Module):\n",
+    "    def __init__(self):\n",
+    "        super().__init__()\n",
+    "        self.scattering = Scattering2D(J=2, shape=(32, 32))\n",
+    "        scatt_dim = 81 * 64\n",
+    "        self.scatt_proj = nn.Linear(scatt_dim, 256)\n",
+    "        self.patchwork = Patchwork(256, num_blocks=2)\n",
+    "\n",
+    "        self.polar_dec = PolarDecompositionFeatures(256, num_features=32)\n",
+    "        self.patchwork_polar = Patchwork(64, num_blocks=2)  # 2*32\n",
+    "        self.head = nn.Linear(64, 10)\n",
+    "\n",
+    "    def forward(self, images):\n",
+    "        scatt = self.scattering(images)\n",
+    "        B = scatt.shape[0]\n",
+    "        scatt_flat = scatt.reshape(B, -1)\n",
+    "        proj = self.scatt_proj(scatt_flat)\n",
+    "        embeddings = l2_normalize(self.patchwork(proj))\n",
+    "\n",
+    "        polar_feat = self.polar_dec(embeddings)\n",
+    "        polar_feat = self.patchwork_polar(polar_feat)\n",
+    "        logits = self.head(polar_feat)\n",
+    "        return logits, embeddings\n",
+    "\n",
+    "model_4_5b = Model_4_5b().to(device)\n",
+    "print(f\"[4.5b] Model created with {sum(p.numel() for p in model_4_5b.parameters() if p.requires_grad):,} params\")\n",
+    "\n",
+    "result_4_5b = run_experiment(\n",
+    "    \"4.5b_polar_decomposition\",\n",
+    "    model_4_5b,\n",
+    "    train_loader,\n",
+    "    val_loader,\n",
+    "    device,\n",
+    "    epochs=100,\n",
+    "    lr=1e-3,\n",
+    "    weight_decay=1e-4,\n",
+    "    patience=20\n",
+    ")\n",
+    "scoreboard.add(result_4_5b)\n",
+    ""
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Final Scoreboard\n",
+    "\n",
+    "All experiments completed. Scoreboard shows results ranked by validation accuracy.\n",
+    ""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "\n",
+    "print(\"\\n\" + \"=\"*80)\n",
+    "print(\"FINAL RESULTS \u2014 GeoLIP Category 4: Invertible Geometric Transforms\")\n",
+    "print(\"=\"*80)\n",
+    "scoreboard._print_table()\n",
+    ""
+   ],
+   "execution_count": null,
+   "outputs": []
+  }
+ ]
+}

spectral/experiment_5_matrix_decompositions.ipynb

@@ -0,0 +1,1118 @@
+{
+ "nbformat": 4,
+ "nbformat_minor": 5,
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.10.0"
+  },
+  "accelerator": "GPU",
+  "colab": {
+   "provenance": [],
+   "gpuType": "T4"
+  }
+ },
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Experiment 5: Matrix Decomposition Signatures\n",
+    "\n",
+    "Decomposing STFT-derived Gram matrices for geometric feature extraction."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "\n",
+    "!pip install -q kymatio scipy tensorboard scikit-learn matplotlib numpy torch torchvision torchaudio tensorboard\n",
+    ""
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Shared Utilities"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "\n",
+    "# \u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\n",
+    "# SHARED UTILITIES \u2014 GeoLIP Hypersphere Experiments\n",
+    "# \u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\n",
+    "import os, sys, time, math, json, datetime, warnings\n",
+    "import numpy as np\n",
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "import torch.nn.functional as F\n",
+    "from torch.utils.data import DataLoader\n",
+    "from torch.utils.tensorboard import SummaryWriter\n",
+    "import torchvision\n",
+    "import torchvision.transforms as T\n",
+    "import matplotlib\n",
+    "matplotlib.use(\"Agg\")\n",
+    "import matplotlib.pyplot as plt\n",
+    "from collections import defaultdict\n",
+    "\n",
+    "warnings.filterwarnings(\"ignore\", category=UserWarning)\n",
+    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
+    "print(f\"[DEVICE] {device}\" + (f\" \u2014 {torch.cuda.get_device_name()}\" if device.type == \"cuda\" else \"\"))\n",
+    "\n",
+    "# \u2500\u2500 CIFAR-10 Loaders \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
+    "CIFAR10_MEAN = (0.4914, 0.4822, 0.4465)\n",
+    "CIFAR10_STD  = (0.2470, 0.2435, 0.2616)\n",
+    "CIFAR10_CLASSES = (\"plane\",\"car\",\"bird\",\"cat\",\"deer\",\"dog\",\"frog\",\"horse\",\"ship\",\"truck\")\n",
+    "\n",
+    "def get_cifar10_loaders(batch_size=256, num_workers=2, augment=True):\n",
+    "    \"\"\"Return train_loader, val_loader for CIFAR-10.\"\"\"\n",
+    "    tf_train = T.Compose([\n",
+    "        T.RandomCrop(32, padding=4),\n",
+    "        T.RandomHorizontalFlip(),\n",
+    "        T.ToTensor(),\n",
+    "        T.Normalize(CIFAR10_MEAN, CIFAR10_STD),\n",
+    "    ]) if augment else T.Compose([T.ToTensor(), T.Normalize(CIFAR10_MEAN, CIFAR10_STD)])\n",
+    "    tf_val = T.Compose([T.ToTensor(), T.Normalize(CIFAR10_MEAN, CIFAR10_STD)])\n",
+    "\n",
+    "    train_ds = torchvision.datasets.CIFAR10(root=\"./data\", train=True,  download=True, transform=tf_train)\n",
+    "    val_ds   = torchvision.datasets.CIFAR10(root=\"./data\", train=False, download=True, transform=tf_val)\n",
+    "    train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True,  num_workers=num_workers, pin_memory=True, drop_last=True)\n",
+    "    val_loader   = DataLoader(val_ds,   batch_size=batch_size, shuffle=False, num_workers=num_workers, pin_memory=True)\n",
+    "    print(f\"[DATA] CIFAR-10: {len(train_ds)} train, {len(val_ds)} val, batch_size={batch_size}\")\n",
+    "    return train_loader, val_loader\n",
+    "\n",
+    "def get_cifar10_raw_loaders(batch_size=256, num_workers=2):\n",
+    "    \"\"\"Return loaders with only ToTensor (no normalization) for custom front-ends.\"\"\"\n",
+    "    tf = T.Compose([T.ToTensor()])\n",
+    "    train_ds = torchvision.datasets.CIFAR10(root=\"./data\", train=True,  download=True, transform=tf)\n",
+    "    val_ds   = torchvision.datasets.CIFAR10(root=\"./data\", train=False, download=True, transform=tf)\n",
+    "    train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True,  num_workers=num_workers, pin_memory=True, drop_last=True)\n",
+    "    val_loader   = DataLoader(val_ds,   batch_size=batch_size, shuffle=False, num_workers=num_workers, pin_memory=True)\n",
+    "    print(f\"[DATA] CIFAR-10 (raw): {len(train_ds)} train, {len(val_ds)} val\")\n",
+    "    return train_loader, val_loader\n",
+    "\n",
+    "# \u2500\u2500 L2 Normalization to Hypersphere \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
+    "def l2_normalize(x, dim=-1, eps=1e-8):\n",
+    "    \"\"\"Project to unit hypersphere S^{d-1}.\"\"\"\n",
+    "    return F.normalize(x, p=2, dim=dim, eps=eps)\n",
+    "\n",
+    "# \u2500\u2500 Constellation Module \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
+    "class Constellation(nn.Module):\n",
+    "    \"\"\"\n",
+    "    K anchors on S^{d-1}. Triangulation: compute distance from input to each anchor.\n",
+    "    Output: K-dim vector of (1 - cos_sim) distances = triangulation profile.\n",
+    "    \"\"\"\n",
+    "    def __init__(self, embed_dim, num_anchors=64, init=\"uniform\"):\n",
+    "        super().__init__()\n",
+    "        self.embed_dim = embed_dim\n",
+    "        self.num_anchors = num_anchors\n",
+    "        anchors = self._init_anchors(embed_dim, num_anchors, init)\n",
+    "        self.anchors = nn.Parameter(anchors, requires_grad=False)  # anchors are NOT learned\n",
+    "\n",
+    "    def _init_anchors(self, d, K, method):\n",
+    "        if method == \"uniform\":\n",
+    "            # QR-based uniform hypersphere initialization\n",
+    "            A = torch.randn(K, d)\n",
+    "            Q, _ = torch.linalg.qr(A.T)\n",
+    "            anchors = Q.T[:K]\n",
+    "            # If K > d, fill remaining with random unit vectors\n",
+    "            if K > d:\n",
+    "                extra = F.normalize(torch.randn(K - d, d), dim=-1)\n",
+    "                anchors = torch.cat([anchors, extra], dim=0)\n",
+    "            return F.normalize(anchors, dim=-1)\n",
+    "        elif method == \"random\":\n",
+    "            return F.normalize(torch.randn(K, d), dim=-1)\n",
+    "        else:\n",
+    "            return F.normalize(torch.randn(K, d), dim=-1)\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        \"\"\"x: (B, d) on S^{d-1} \u2192 (B, K) triangulation distances.\"\"\"\n",
+    "        x_norm = l2_normalize(x)\n",
+    "        cos_sim = torch.mm(x_norm, self.anchors.T)  # (B, K)\n",
+    "        distances = 1.0 - cos_sim  # in [0, 2]\n",
+    "        return distances\n",
+    "\n",
+    "    def push_anchors(self, embeddings, labels, num_classes=10, momentum=0.1):\n",
+    "        \"\"\"Push anchors toward class centroids (call periodically).\"\"\"\n",
+    "        with torch.no_grad():\n",
+    "            for c in range(num_classes):\n",
+    "                mask = (labels == c)\n",
+    "                if mask.sum() < 2:\n",
+    "                    continue\n",
+    "                centroid = l2_normalize(embeddings[mask].mean(dim=0, keepdim=True))\n",
+    "                # Assign anchors to this class based on nearest centroid\n",
+    "                start = c * (self.num_anchors // num_classes)\n",
+    "                end = start + (self.num_anchors // num_classes)\n",
+    "                self.anchors.data[start:end] = l2_normalize(\n",
+    "                    (1 - momentum) * self.anchors.data[start:end] + momentum * centroid\n",
+    "                )\n",
+    "\n",
+    "    def anchor_activation(self, x):\n",
+    "        \"\"\"Return fraction of anchors that are 'active' (nearest to at least one sample).\"\"\"\n",
+    "        x_norm = l2_normalize(x)\n",
+    "        cos_sim = torch.mm(x_norm, self.anchors.T)\n",
+    "        nearest = cos_sim.argmax(dim=1)\n",
+    "        active = len(nearest.unique())\n",
+    "        return active, self.num_anchors\n",
+    "\n",
+    "# \u2500\u2500 Patchwork Block \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
+    "class SquaredReLU(nn.Module):\n",
+    "    def forward(self, x):\n",
+    "        return F.relu(x) ** 2\n",
+    "\n",
+    "class PatchworkBlock(nn.Module):\n",
+    "    \"\"\"SquaredReLU + LayerNorm block for post-constellation processing.\"\"\"\n",
+    "    def __init__(self, dim):\n",
+    "        super().__init__()\n",
+    "        self.net = nn.Sequential(\n",
+    "            nn.Linear(dim, dim),\n",
+    "            SquaredReLU(),\n",
+    "            nn.LayerNorm(dim),\n",
+    "            nn.Linear(dim, dim),\n",
+    "            SquaredReLU(),\n",
+    "            nn.LayerNorm(dim),\n",
+    "        )\n",
+    "    def forward(self, x):\n",
+    "        return x + self.net(x)  # residual\n",
+    "\n",
+    "class Patchwork(nn.Module):\n",
+    "    \"\"\"Stack of PatchworkBlocks.\"\"\"\n",
+    "    def __init__(self, dim, num_blocks=2):\n",
+    "        super().__init__()\n",
+    "        self.blocks = nn.Sequential(*[PatchworkBlock(dim) for _ in range(num_blocks)])\n",
+    "    def forward(self, x):\n",
+    "        return self.blocks(x)\n",
+    "\n",
+    "# \u2500\u2500 Classifier Head \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
+    "class ConstellationClassifier(nn.Module):\n",
+    "    \"\"\"Full pipeline: encoder output \u2192 constellation \u2192 patchwork \u2192 classify.\"\"\"\n",
+    "    def __init__(self, embed_dim, num_anchors=64, num_classes=10, patchwork_blocks=2):\n",
+    "        super().__init__()\n",
+    "        self.constellation = Constellation(embed_dim, num_anchors)\n",
+    "        self.patchwork = Patchwork(num_anchors, patchwork_blocks)\n",
+    "        self.head = nn.Linear(num_anchors, num_classes)\n",
+    "\n",
+    "    def forward(self, embeddings):\n",
+    "        tri = self.constellation(embeddings)       # (B, K)\n",
+    "        pw  = self.patchwork(tri)                   # (B, K)\n",
+    "        logits = self.head(pw)                      # (B, 10)\n",
+    "        return logits, tri\n",
+    "\n",
+    "    def count_params(self):\n",
+    "        return sum(p.numel() for p in self.parameters() if p.requires_grad)\n",
+    "\n",
+    "# \u2500\u2500 Losses \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
+    "class CayleyMengerVolumeLoss(nn.Module):\n",
+    "    \"\"\"CV loss: encourage constellation volume in 0.18-0.25 band.\"\"\"\n",
+    "    def __init__(self, target_low=0.18, target_high=0.25):\n",
+    "        super().__init__()\n",
+    "        self.target_low = target_low\n",
+    "        self.target_high = target_high\n",
+    "\n",
+    "    def forward(self, tri):\n",
+    "        \"\"\"tri: (B, K) triangulation distances.\"\"\"\n",
+    "        # Approximate volume via determinant of distance submatrix\n",
+    "        B, K = tri.shape\n",
+    "        # Use std of triangulation as proxy for volume spread\n",
+    "        vol = tri.std(dim=1).mean()\n",
+    "        # Penalty if outside target band\n",
+    "        below = F.relu(self.target_low - vol)\n",
+    "        above = F.relu(vol - self.target_high)\n",
+    "        return below + above, vol.item()\n",
+    "\n",
+    "class InfoNCELoss(nn.Module):\n",
+    "    \"\"\"Self-supervised contrastive loss on augmented views.\"\"\"\n",
+    "    def __init__(self, temperature=0.07):\n",
+    "        super().__init__()\n",
+    "        self.temperature = temperature\n",
+    "\n",
+    "    def forward(self, z1, z2):\n",
+    "        \"\"\"z1, z2: (B, d) L2-normalized embeddings of two augmented views.\"\"\"\n",
+    "        z1 = l2_normalize(z1)\n",
+    "        z2 = l2_normalize(z2)\n",
+    "        B = z1.shape[0]\n",
+    "        z = torch.cat([z1, z2], dim=0)  # (2B, d)\n",
+    "        sim = torch.mm(z, z.T) / self.temperature  # (2B, 2B)\n",
+    "        # Mask out self-similarity\n",
+    "        mask = ~torch.eye(2*B, dtype=bool, device=z.device)\n",
+    "        sim = sim.masked_fill(~mask, -1e9)\n",
+    "        # Positive pairs: (i, i+B) and (i+B, i)\n",
+    "        labels = torch.cat([torch.arange(B, 2*B), torch.arange(B)]).to(z.device)\n",
+    "        loss = F.cross_entropy(sim, labels)\n",
+    "        return loss\n",
+    "\n",
+    "class ClassConditionalInfoNCE(nn.Module):\n",
+    "    \"\"\"Class-conditional contrastive: positive = same class, not same image.\"\"\"\n",
+    "    def __init__(self, temperature=0.07):\n",
+    "        super().__init__()\n",
+    "        self.temperature = temperature\n",
+    "\n",
+    "    def forward(self, z, labels):\n",
+    "        z = l2_normalize(z)\n",
+    "        B = z.shape[0]\n",
+    "        sim = torch.mm(z, z.T) / self.temperature\n",
+    "        # Mask: positive if same class, negative otherwise\n",
+    "        label_mask = (labels.unsqueeze(0) == labels.unsqueeze(1)).float()\n",
+    "        label_mask.fill_diagonal_(0)  # exclude self\n",
+    "        # For each sample, compute log-softmax over all non-self, weighted by positives\n",
+    "        logits = sim - torch.diag(torch.full((B,), 1e9, device=z.device))\n",
+    "        pos_mask = label_mask\n",
+    "        neg_mask = 1 - label_mask\n",
+    "        neg_mask.fill_diagonal_(0)\n",
+    "        # InfoNCE: -log(sum_pos exp(sim) / sum_all exp(sim))\n",
+    "        exp_sim = torch.exp(logits)\n",
+    "        pos_sum = (exp_sim * pos_mask).sum(dim=1)\n",
+    "        all_sum = (exp_sim * (pos_mask + neg_mask)).sum(dim=1)\n",
+    "        loss = -torch.log(pos_sum / (all_sum + 1e-8) + 1e-8).mean()\n",
+    "        return loss\n",
+    "\n",
+    "# \u2500\u2500 TensorBoard Helpers \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
+    "def make_writer(experiment_id, base_dir=\"runs\"):\n",
+    "    \"\"\"Create a TensorBoard SummaryWriter with timestamped run directory.\"\"\"\n",
+    "    timestamp = datetime.datetime.now().strftime(\"%Y%m%d_%H%M%S\")\n",
+    "    log_dir = os.path.join(base_dir, f\"{experiment_id}_{timestamp}\")\n",
+    "    writer = SummaryWriter(log_dir=log_dir)\n",
+    "    print(f\"[TB] Logging to {log_dir}\")\n",
+    "    return writer\n",
+    "\n",
+    "def log_epoch_metrics(writer, epoch, metrics_dict, prefix=\"\"):\n",
+    "    \"\"\"Log a dict of metrics to TensorBoard.\"\"\"\n",
+    "    for key, val in metrics_dict.items():\n",
+    "        tag = f\"{prefix}/{key}\" if prefix else key\n",
+    "        writer.add_scalar(tag, val, epoch)\n",
+    "\n",
+    "def log_embedding_viz(writer, embeddings, labels, epoch, tag=\"embeddings\"):\n",
+    "    \"\"\"Log embeddings for TensorBoard projector.\"\"\"\n",
+    "    meta = [CIFAR10_CLASSES[l] for l in labels]\n",
+    "    writer.add_embedding(embeddings, metadata=meta, global_step=epoch, tag=tag)\n",
+    "\n",
+    "def log_anchor_histogram(writer, constellation, epoch):\n",
+    "    \"\"\"Log anchor norm distribution.\"\"\"\n",
+    "    norms = constellation.anchors.data.norm(dim=1)\n",
+    "    writer.add_histogram(\"anchors/norms\", norms, epoch)\n",
+    "    cos_sim_matrix = torch.mm(constellation.anchors.data, constellation.anchors.data.T)\n",
+    "    off_diag = cos_sim_matrix[~torch.eye(cos_sim_matrix.shape[0], dtype=bool)].flatten()\n",
+    "    writer.add_histogram(\"anchors/inter_cosine_sim\", off_diag, epoch)\n",
+    "\n",
+    "def log_confusion_matrix(writer, all_preds, all_labels, epoch):\n",
+    "    \"\"\"Log confusion matrix as image to TensorBoard.\"\"\"\n",
+    "    from sklearn.metrics import confusion_matrix\n",
+    "    cm = confusion_matrix(all_labels, all_preds, labels=range(10))\n",
+    "    fig, ax = plt.subplots(1, 1, figsize=(8, 8))\n",
+    "    im = ax.imshow(cm, interpolation=\"nearest\", cmap=plt.cm.Blues)\n",
+    "    ax.set_title(f\"Confusion Matrix \u2014 Epoch {epoch}\")\n",
+    "    ax.set_xlabel(\"Predicted\")\n",
+    "    ax.set_ylabel(\"True\")\n",
+    "    tick_marks = range(10)\n",
+    "    ax.set_xticks(tick_marks)\n",
+    "    ax.set_xticklabels(CIFAR10_CLASSES, rotation=45, ha=\"right\")\n",
+    "    ax.set_yticks(tick_marks)\n",
+    "    ax.set_yticklabels(CIFAR10_CLASSES)\n",
+    "    for i in range(10):\n",
+    "        for j in range(10):\n",
+    "            ax.text(j, i, str(cm[i, j]), ha=\"center\", va=\"center\",\n",
+    "                    color=\"white\" if cm[i, j] > cm.max() / 2 else \"black\", fontsize=7)\n",
+    "    fig.tight_layout()\n",
+    "    writer.add_figure(\"confusion_matrix\", fig, epoch)\n",
+    "    plt.close(fig)\n",
+    "\n",
+    "def log_triangulation_stats(writer, tri, epoch):\n",
+    "    \"\"\"Log triangulation distance statistics.\"\"\"\n",
+    "    writer.add_scalar(\"tri/mean\", tri.mean().item(), epoch)\n",
+    "    writer.add_scalar(\"tri/std\", tri.std().item(), epoch)\n",
+    "    writer.add_scalar(\"tri/min\", tri.min().item(), epoch)\n",
+    "    writer.add_scalar(\"tri/max\", tri.max().item(), epoch)\n",
+    "    writer.add_histogram(\"tri/distances\", tri.detach().cpu(), epoch)\n",
+    "\n",
+    "def log_scatter_2d(writer, embeddings, labels, epoch, tag=\"scatter_2d\", method=\"pca\"):\n",
+    "    \"\"\"Log 2D scatter plot of embeddings projected via PCA or t-SNE.\"\"\"\n",
+    "    emb_np = embeddings.detach().cpu().numpy()\n",
+    "    lab_np = labels.detach().cpu().numpy() if torch.is_tensor(labels) else np.array(labels)\n",
+    "    if method == \"pca\":\n",
+    "        from sklearn.decomposition import PCA\n",
+    "        proj = PCA(n_components=2).fit_transform(emb_np)\n",
+    "    else:\n",
+    "        from sklearn.manifold import TSNE\n",
+    "        proj = TSNE(n_components=2, perplexity=30, n_iter=500).fit_transform(emb_np)\n",
+    "    fig, ax = plt.subplots(1, 1, figsize=(8, 8))\n",
+    "    for c in range(10):\n",
+    "        mask = lab_np == c\n",
+    "        ax.scatter(proj[mask, 0], proj[mask, 1], s=3, alpha=0.5, label=CIFAR10_CLASSES[c])\n",
+    "    ax.legend(markerscale=5, fontsize=8)\n",
+    "    ax.set_title(f\"{tag} ({method.upper()}) \u2014 Epoch {epoch}\")\n",
+    "    fig.tight_layout()\n",
+    "    writer.add_figure(f\"{tag}/{method}\", fig, epoch)\n",
+    "    plt.close(fig)\n",
+    "\n",
+    "# \u2500\u2500 Training Utilities \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
+    "def train_one_epoch(model, loader, optimizer, criterion, device, epoch=0,\n",
+    "                    scheduler=None, extra_losses=None):\n",
+    "    \"\"\"\n",
+    "    Generic training loop for one epoch.\n",
+    "    model: callable(batch) \u2192 (logits, embeddings_on_sphere)\n",
+    "    extra_losses: list of (name, loss_fn, weight) where loss_fn(embeddings) \u2192 scalar\n",
+    "    \"\"\"\n",
+    "    model.train()\n",
+    "    total_loss = 0\n",
+    "    total_ce = 0\n",
+    "    correct = 0\n",
+    "    total = 0\n",
+    "    loss_components = defaultdict(float)\n",
+    "\n",
+    "    for batch_idx, (images, labels) in enumerate(loader):\n",
+    "        images, labels = images.to(device), labels.to(device)\n",
+    "        optimizer.zero_grad()\n",
+    "\n",
+    "        logits, embeddings = model(images)\n",
+    "        ce_loss = criterion(logits, labels)\n",
+    "        loss = ce_loss\n",
+    "        loss_components[\"CE\"] += ce_loss.item()\n",
+    "\n",
+    "        if extra_losses:\n",
+    "            for name, loss_fn, weight in extra_losses:\n",
+    "                if name == \"CV\":\n",
+    "                    el, vol_val = loss_fn(embeddings)\n",
+    "                    loss_components[\"CV_vol\"] += vol_val\n",
+    "                elif name == \"InfoNCE_class\":\n",
+    "                    el = loss_fn(embeddings, labels)\n",
+    "                else:\n",
+    "                    el = loss_fn(embeddings)\n",
+    "                loss = loss + weight * el\n",
+    "                loss_components[name] += el.item()\n",
+    "\n",
+    "        loss.backward()\n",
+    "        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)\n",
+    "        optimizer.step()\n",
+    "        if scheduler:\n",
+    "            scheduler.step()\n",
+    "\n",
+    "        total_loss += loss.item()\n",
+    "        _, predicted = logits.max(1)\n",
+    "        total += labels.size(0)\n",
+    "        correct += predicted.eq(labels).sum().item()\n",
+    "\n",
+    "    n_batches = len(loader)\n",
+    "    metrics = {\n",
+    "        \"loss/total\": total_loss / n_batches,\n",
+    "        \"loss/CE\": loss_components[\"CE\"] / n_batches,\n",
+    "        \"accuracy/train\": 100.0 * correct / total,\n",
+    "    }\n",
+    "    for k, v in loss_components.items():\n",
+    "        if k not in (\"CE\",):\n",
+    "            metrics[f\"loss/{k}\"] = v / n_batches\n",
+    "    return metrics\n",
+    "\n",
+    "@torch.no_grad()\n",
+    "def evaluate(model, loader, criterion, device, collect_embeddings=False):\n",
+    "    \"\"\"Evaluate model on validation set. Optionally collect embeddings.\"\"\"\n",
+    "    model.eval()\n",
+    "    total_loss = 0\n",
+    "    correct = 0\n",
+    "    total = 0\n",
+    "    all_preds = []\n",
+    "    all_labels = []\n",
+    "    all_embeddings = []\n",
+    "\n",
+    "    for images, labels in loader:\n",
+    "        images, labels = images.to(device), labels.to(device)\n",
+    "        logits, embeddings = model(images)\n",
+    "        loss = criterion(logits, labels)\n",
+    "        total_loss += loss.item()\n",
+    "        _, predicted = logits.max(1)\n",
+    "        total += labels.size(0)\n",
+    "        correct += predicted.eq(labels).sum().item()\n",
+    "        all_preds.extend(predicted.cpu().numpy())\n",
+    "        all_labels.extend(labels.cpu().numpy())\n",
+    "        if collect_embeddings:\n",
+    "            all_embeddings.append(embeddings.cpu())\n",
+    "\n",
+    "    n_batches = len(loader)\n",
+    "    metrics = {\n",
+    "        \"loss/val\": total_loss / n_batches,\n",
+    "        \"accuracy/val\": 100.0 * correct / total,\n",
+    "    }\n",
+    "    result = {\"metrics\": metrics, \"preds\": np.array(all_preds), \"labels\": np.array(all_labels)}\n",
+    "    if collect_embeddings:\n",
+    "        result[\"embeddings\"] = torch.cat(all_embeddings, dim=0)\n",
+    "    return result\n",
+    "\n",
+    "def run_experiment(experiment_id, model, train_loader, val_loader, device,\n",
+    "                   epochs=100, lr=1e-3, weight_decay=1e-4, extra_losses=None,\n",
+    "                   push_anchors_every=50, log_embeddings_every=25,\n",
+    "                   patience=20):\n",
+    "    \"\"\"\n",
+    "    Full experiment runner with TensorBoard logging, early stopping, and diagnostics.\n",
+    "    model must return (logits, sphere_embeddings).\n",
+    "    \"\"\"\n",
+    "    writer = make_writer(experiment_id)\n",
+    "    criterion = nn.CrossEntropyLoss()\n",
+    "    optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=weight_decay)\n",
+    "    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs)\n",
+    "\n",
+    "    best_val_acc = 0.0\n",
+    "    best_epoch = 0\n",
+    "    no_improve = 0\n",
+    "\n",
+    "    # Log model info\n",
+    "    n_params = sum(p.numel() for p in model.parameters() if p.requires_grad)\n",
+    "    writer.add_text(\"config/params\", f\"Learned params: {n_params:,}\")\n",
+    "    writer.add_text(\"config/experiment\", experiment_id)\n",
+    "    print(f\"\\n{'='*70}\")\n",
+    "    print(f\"[EXP] {experiment_id}\")\n",
+    "    print(f\"[EXP] Learned params: {n_params:,}, Epochs: {epochs}, LR: {lr}\")\n",
+    "    print(f\"{'='*70}\")\n",
+    "\n",
+    "    for epoch in range(1, epochs + 1):\n",
+    "        t0 = time.time()\n",
+    "\n",
+    "        # Train\n",
+    "        train_metrics = train_one_epoch(\n",
+    "            model, train_loader, optimizer, criterion, device, epoch,\n",
+    "            extra_losses=extra_losses\n",
+    "        )\n",
+    "\n",
+    "        # Evaluate\n",
+    "        collect_emb = (epoch % log_embeddings_every == 0) or (epoch == epochs)\n",
+    "        val_result = evaluate(model, val_loader, criterion, device, collect_embeddings=collect_emb)\n",
+    "        val_metrics = val_result[\"metrics\"]\n",
+    "\n",
+    "        # Combine metrics\n",
+    "        all_metrics = {**train_metrics, **val_metrics}\n",
+    "        all_metrics[\"lr\"] = optimizer.param_groups[0][\"lr\"]\n",
+    "        all_metrics[\"epoch_time\"] = time.time() - t0\n",
+    "        all_metrics[\"train_val_gap\"] = train_metrics[\"accuracy/train\"] - val_metrics[\"accuracy/val\"]\n",
+    "\n",
+    "        # Constellation diagnostics\n",
+    "        if hasattr(model, \"constellation\"):\n",
+    "            constellation = model.constellation\n",
+    "        elif hasattr(model, \"classifier\") and hasattr(model.classifier, \"constellation\"):\n",
+    "            constellation = model.classifier.constellation\n",
+    "        else:\n",
+    "            constellation = None\n",
+    "\n",
+    "        if constellation is not None:\n",
+    "            log_anchor_histogram(writer, constellation, epoch)\n",
+    "\n",
+    "            # Push anchors periodically\n",
+    "            if epoch % push_anchors_every == 0 and \"embeddings\" in val_result:\n",
+    "                emb = val_result[\"embeddings\"][:2000].to(device)\n",
+    "                lab = torch.tensor(val_result[\"labels\"][:2000]).to(device)\n",
+    "                constellation.push_anchors(emb, lab)\n",
+    "                active, total_anch = constellation.anchor_activation(emb)\n",
+    "                all_metrics[\"anchors/active\"] = active\n",
+    "                all_metrics[\"anchors/total\"] = total_anch\n",
+    "                all_metrics[\"anchors/utilization\"] = active / total_anch\n",
+    "                print(f\"  [ANCHOR PUSH] Active: {active}/{total_anch} ({100*active/total_anch:.1f}%)\")\n",
+    "\n",
+    "        # Log all metrics\n",
+    "        log_epoch_metrics(writer, epoch, all_metrics)\n",
+    "        scheduler.step()\n",
+    "\n",
+    "        # Log embeddings visualization\n",
+    "        if collect_emb and \"embeddings\" in val_result:\n",
+    "            emb = val_result[\"embeddings\"][:2000]\n",
+    "            lab = val_result[\"labels\"][:2000]\n",
+    "            log_scatter_2d(writer, emb, lab, epoch, tag=\"val_embeddings\", method=\"pca\")\n",
+    "            if epoch == epochs:\n",
+    "                log_confusion_matrix(writer, val_result[\"preds\"], val_result[\"labels\"], epoch)\n",
+    "                log_scatter_2d(writer, emb, lab, epoch, tag=\"val_embeddings_tsne\", method=\"tsne\")\n",
+    "\n",
+    "        # Log triangulation stats from a sample batch\n",
+    "        if constellation is not None and collect_emb and \"embeddings\" in val_result:\n",
+    "            sample_emb = val_result[\"embeddings\"][:500].to(device)\n",
+    "            tri_sample = constellation(sample_emb)\n",
+    "            log_triangulation_stats(writer, tri_sample, epoch)\n",
+    "\n",
+    "        # Best model tracking\n",
+    "        val_acc = val_metrics[\"accuracy/val\"]\n",
+    "        if val_acc > best_val_acc:\n",
+    "            best_val_acc = val_acc\n",
+    "            best_epoch = epoch\n",
+    "            no_improve = 0\n",
+    "            writer.add_scalar(\"best/val_acc\", best_val_acc, epoch)\n",
+    "            writer.add_scalar(\"best/epoch\", best_epoch, epoch)\n",
+    "        else:\n",
+    "            no_improve += 1\n",
+    "\n",
+    "        # Print progress\n",
+    "        if epoch % 5 == 0 or epoch == 1 or epoch == epochs:\n",
+    "            print(f\"  E{epoch:3d} | TrainAcc {train_metrics['accuracy/train']:5.1f}% | \"\n",
+    "                  f\"ValAcc {val_acc:5.1f}% | Loss {all_metrics['loss/total']:.4f} | \"\n",
+    "                  f\"Best {best_val_acc:.1f}% @E{best_epoch} | \"\n",
+    "                  f\"Gap {all_metrics['train_val_gap']:+.1f}% | \"\n",
+    "                  f\"{all_metrics['epoch_time']:.1f}s\")\n",
+    "\n",
+    "        # Early stopping\n",
+    "        if no_improve >= patience:\n",
+    "            print(f\"  [EARLY STOP] No improvement for {patience} epochs. Best: {best_val_acc:.1f}% @E{best_epoch}\")\n",
+    "            break\n",
+    "\n",
+    "    writer.close()\n",
+    "    print(f\"\\n[RESULT] {experiment_id}: Best Val Acc = {best_val_acc:.2f}% at epoch {best_epoch}\")\n",
+    "    print(f\"[RESULT] Learned params: {n_params:,}\")\n",
+    "    return {\"experiment\": experiment_id, \"best_val_acc\": best_val_acc, \"best_epoch\": best_epoch,\n",
+    "            \"params\": n_params, \"final_train_acc\": train_metrics[\"accuracy/train\"]}\n",
+    "\n",
+    "# \u2500\u2500 Johnson-Lindenstrauss Projection \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
+    "class JLProjection(nn.Module):\n",
+    "    \"\"\"Fixed random Gaussian projection (not learned).\"\"\"\n",
+    "    def __init__(self, input_dim, output_dim):\n",
+    "        super().__init__()\n",
+    "        self.register_buffer(\"W\", torch.randn(input_dim, output_dim) / math.sqrt(output_dim))\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        return x @ self.W\n",
+    "\n",
+    "# \u2500\u2500 Results Scoreboard \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
+    "class Scoreboard:\n",
+    "    \"\"\"Track results across experiments in a session.\"\"\"\n",
+    "    def __init__(self):\n",
+    "        self.results = []\n",
+    "\n",
+    "    def add(self, result_dict):\n",
+    "        self.results.append(result_dict)\n",
+    "        self._print_table()\n",
+    "\n",
+    "    def _print_table(self):\n",
+    "        print(f\"\\n{'='*80}\")\n",
+    "        print(f\"{'SCOREBOARD':^80}\")\n",
+    "        print(f\"{'='*80}\")\n",
+    "        print(f\"{'Experiment':<40} {'Val Acc':>8} {'Params':>10} {'Best E':>7}\")\n",
+    "        print(f\"{'-'*40} {'-'*8} {'-'*10} {'-'*7}\")\n",
+    "        for r in sorted(self.results, key=lambda x: x[\"best_val_acc\"], reverse=True):\n",
+    "            print(f\"{r['experiment']:<40} {r['best_val_acc']:>7.2f}% {r['params']:>10,} {r['best_epoch']:>7}\")\n",
+    "        print(f\"{'='*80}\\n\")\n",
+    "\n",
+    "    def save(self, path=\"scoreboard.json\"):\n",
+    "        with open(path, \"w\") as f:\n",
+    "            json.dump(self.results, f, indent=2)\n",
+    "        print(f\"[SCOREBOARD] Saved to {path}\")\n",
+    "\n",
+    "scoreboard = Scoreboard()\n",
+    "print(\"[SHARED] All utilities loaded successfully.\")\n",
+    ""
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 5.2a: QR Decomposition Signatures\n",
+    "\n",
+    "Rotation matrix Q + diagonal scale R \u2192 features"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "\n",
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "import torch.nn.functional as F\n",
+    "import numpy as np\n",
+    "\n",
+    "print(\"[5.2a] QR Decomposition Signatures\")\n",
+    "print(\"=\"*70)\n",
+    "\n",
+    "class STFTFrontend(nn.Module):\n",
+    "    '''STFT front-end: row-wise + col-wise STFT on each channel.'''\n",
+    "    def __init__(self, n_fft=16, hop_length=4, freq_bins=8):\n",
+    "        super().__init__()\n",
+    "        self.n_fft = n_fft\n",
+    "        self.hop_length = hop_length\n",
+    "        self.freq_bins = freq_bins\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        # x: (B, 3, 32, 32) \u2192 (B, 24, 32, 32) [8 freq channels per color]\n",
+    "        B, C, H, W = x.shape\n",
+    "        stft_features = []\n",
+    "\n",
+    "        for c in range(C):\n",
+    "            channel = x[:, c:c+1, :, :]  # (B, 1, 32, 32)\n",
+    "\n",
+    "            # Row-wise STFT\n",
+    "            rows = channel.view(B * H, W)  # (B*32, 32)\n",
+    "            row_stft = torch.stft(rows, self.n_fft, self.hop_length, return_complex=True,\n",
+    "                                  pad_mode='reflect', window=torch.hann_window(self.n_fft, device=x.device))\n",
+    "            row_mag = torch.abs(row_stft)[:, :self.freq_bins, :]  # (B*32, 8, ~6)\n",
+    "            row_mag = row_mag.mean(dim=2)  # (B*32, 8)\n",
+    "            row_mag = row_mag.view(B, H, self.freq_bins)  # (B, 32, 8)\n",
+    "\n",
+    "            # Col-wise STFT\n",
+    "            cols = channel.view(B * W, H)  # (B*32, 32)\n",
+    "            col_stft = torch.stft(cols, self.n_fft, self.hop_length, return_complex=True,\n",
+    "                                  pad_mode='reflect', window=torch.hann_window(self.n_fft, device=x.device))\n",
+    "            col_mag = torch.abs(col_stft)[:, :self.freq_bins, :]  # (B*32, 8, ~6)\n",
+    "            col_mag = col_mag.mean(dim=2)  # (B*32, 8)\n",
+    "            col_mag = col_mag.view(B, W, self.freq_bins)  # (B, 32, 8)\n",
+    "\n",
+    "            # Reshape to image-like (B, 1, spatial_freq_combined)\n",
+    "            combined = torch.cat([row_mag.view(B, -1), col_mag.view(B, -1)], dim=1)  # (B, 512)\n",
+    "            stft_features.append(combined)\n",
+    "\n",
+    "        return torch.cat(stft_features, dim=1)  # (B, 1536)\n",
+    "\n",
+    "class QRDecompositionModel(nn.Module):\n",
+    "    '''STFT \u2192 Gram matrix \u2192 QR decomposition \u2192 features.'''\n",
+    "    def __init__(self, embed_dim=128, num_classes=10):\n",
+    "        super().__init__()\n",
+    "        self.stft = STFTFrontend()\n",
+    "        self.dim_reduce = nn.Linear(1536, 256)  # Reduce STFT features\n",
+    "        self.gram_dim = 256\n",
+    "\n",
+    "        # QR-extracted features\n",
+    "        self.feat_dim = self.gram_dim + self.gram_dim  # Q trace + R diagonal\n",
+    "        self.embed = nn.Linear(self.feat_dim, embed_dim)\n",
+    "        self.classifier = ConstellationClassifier(embed_dim, num_anchors=64, num_classes=num_classes)\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        # STFT features\n",
+    "        stft_feat = self.stft(x)  # (B, 1536)\n",
+    "        feat_red = self.dim_reduce(stft_feat)  # (B, 256)\n",
+    "\n",
+    "        # Form Gram matrix\n",
+    "        B = feat_red.shape[0]\n",
+    "        gram = torch.bmm(feat_red.unsqueeze(2), feat_red.unsqueeze(1))  # (B, 256, 256)\n",
+    "\n",
+    "        # QR decomposition\n",
+    "        Q, R = torch.linalg.qr(gram)\n",
+    "        Q_trace = torch.trace(Q.view(-1, self.gram_dim, self.gram_dim).transpose(-1, -2)).view(B)\n",
+    "        R_diag = torch.diagonal(R, dim1=1, dim2=2)  # (B, 256)\n",
+    "\n",
+    "        # Extract features\n",
+    "        qr_feat = torch.cat([Q_trace.unsqueeze(1), R_diag.mean(dim=1, keepdim=True)], dim=1)  # (B, 2)\n",
+    "\n",
+    "        # Combine with reduced features\n",
+    "        combined = torch.cat([feat_red, qr_feat], dim=1)  # (B, 258)\n",
+    "\n",
+    "        # Embed to sphere\n",
+    "        emb = l2_normalize(self.embed(combined))  # (B, 128)\n",
+    "\n",
+    "        # Classify\n",
+    "        logits, tri = self.classifier(emb)\n",
+    "        return logits, emb\n",
+    "\n",
+    "model = QRDecompositionModel(embed_dim=128, num_classes=10).to(device)\n",
+    "print(f\"[5.2a] Model params: {sum(p.numel() for p in model.parameters()):,}\")\n",
+    "\n",
+    "train_loader, val_loader = get_cifar10_raw_loaders(batch_size=256)\n",
+    "\n",
+    "result_5_2a = run_experiment(\n",
+    "    \"5.2a_QR_Decomposition\",\n",
+    "    model, train_loader, val_loader, device,\n",
+    "    epochs=100, lr=1e-3, weight_decay=1e-4, patience=20\n",
+    ")\n",
+    "scoreboard.add(result_5_2a)\n",
+    ""
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 5.2b: Schur Decomposition Signatures\n",
+    "\n",
+    "T diagonal (eigenvalues) + T off-diagonal (coupling) \u2192 features"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "\n",
+    "print(\"[5.2b] Schur Decomposition Signatures\")\n",
+    "print(\"=\"*70)\n",
+    "\n",
+    "class SchurDecompositionModel(nn.Module):\n",
+    "    '''STFT \u2192 Gram matrix \u2192 Schur decomposition.'''\n",
+    "    def __init__(self, embed_dim=128, num_classes=10):\n",
+    "        super().__init__()\n",
+    "        self.stft = STFTFrontend()\n",
+    "        self.dim_reduce = nn.Linear(1536, 256)\n",
+    "        self.gram_dim = 256\n",
+    "\n",
+    "        self.feat_dim = self.gram_dim + 10  # Eigenvalue stats + coupling\n",
+    "        self.embed = nn.Linear(self.feat_dim, embed_dim)\n",
+    "        self.classifier = ConstellationClassifier(embed_dim, num_anchors=64, num_classes=num_classes)\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        stft_feat = self.stft(x)\n",
+    "        feat_red = self.dim_reduce(stft_feat)  # (B, 256)\n",
+    "\n",
+    "        # Gram matrix\n",
+    "        B = feat_red.shape[0]\n",
+    "        gram = torch.bmm(feat_red.unsqueeze(2), feat_red.unsqueeze(1))  # (B, 256, 256)\n",
+    "\n",
+    "        # Schur decomposition (similar to eigendecomposition for symmetric)\n",
+    "        eigenvals, eigenvecs = torch.linalg.eigh(gram)\n",
+    "\n",
+    "        # Extract features from eigenvalues\n",
+    "        eig_mean = eigenvals.mean(dim=1)\n",
+    "        eig_std = eigenvals.std(dim=1)\n",
+    "        eig_max = eigenvals.max(dim=1)[0]\n",
+    "        eig_min = eigenvals.min(dim=1)[0]\n",
+    "        eig_log_sum = torch.log(torch.clamp(eigenvals, min=1e-8)).sum(dim=1)\n",
+    "\n",
+    "        # Off-diagonal coupling (correlation of eigenvectors)\n",
+    "        offdiag = gram - torch.diag_embed(torch.diagonal(gram, dim1=1, dim2=2))\n",
+    "        coupling = offdiag.abs().sum(dim=(1, 2))\n",
+    "\n",
+    "        # Combine features\n",
+    "        schur_feat = torch.stack([eig_mean, eig_std, eig_max, eig_min, eig_log_sum, coupling,\n",
+    "                                   coupling/self.gram_dim, eigenvals[:, 0], eigenvals[:, -1],\n",
+    "                                   eigenvals.sum(dim=1)], dim=1)  # (B, 10)\n",
+    "\n",
+    "        combined = torch.cat([feat_red, schur_feat], dim=1)\n",
+    "\n",
+    "        emb = l2_normalize(self.embed(combined))\n",
+    "        logits, tri = self.classifier(emb)\n",
+    "        return logits, emb\n",
+    "\n",
+    "model = SchurDecompositionModel(embed_dim=128, num_classes=10).to(device)\n",
+    "print(f\"[5.2b] Model params: {sum(p.numel() for p in model.parameters()):,}\")\n",
+    "\n",
+    "result_5_2b = run_experiment(\n",
+    "    \"5.2b_Schur_Decomposition\",\n",
+    "    model, train_loader, val_loader, device,\n",
+    "    epochs=100, lr=1e-3, weight_decay=1e-4, patience=20\n",
+    ")\n",
+    "scoreboard.add(result_5_2b)\n",
+    ""
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 5.2c: Eigendecomposition with Signature Comparison\n",
+    "\n",
+    "Eigenvalues vs SVD singular values vs Cholesky diagonal"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "\n",
+    "print(\"[5.2c] Eigendecomposition of Gram & Signature Comparison\")\n",
+    "print(\"=\"*70)\n",
+    "\n",
+    "class EigenSignatureModel(nn.Module):\n",
+    "    '''Compare eigenvalues (Gram eigh), SVD (Gram svd), Cholesky (L diag).'''\n",
+    "    def __init__(self, embed_dim=128, num_classes=10):\n",
+    "        super().__init__()\n",
+    "        self.stft = STFTFrontend()\n",
+    "        self.dim_reduce = nn.Linear(1536, 256)\n",
+    "        self.gram_dim = 256\n",
+    "\n",
+    "        # Three signature types \u00d7 10 stats each\n",
+    "        self.feat_dim = 256 + 30\n",
+    "        self.embed = nn.Linear(self.feat_dim, embed_dim)\n",
+    "        self.classifier = ConstellationClassifier(embed_dim, num_anchors=64, num_classes=num_classes)\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        stft_feat = self.stft(x)\n",
+    "        feat_red = self.dim_reduce(stft_feat)\n",
+    "\n",
+    "        B = feat_red.shape[0]\n",
+    "        gram = torch.bmm(feat_red.unsqueeze(2), feat_red.unsqueeze(1))\n",
+    "\n",
+    "        # Add small diagonal to ensure PSD\n",
+    "        gram = gram + torch.eye(self.gram_dim, device=gram.device).unsqueeze(0) * 1e-6\n",
+    "\n",
+    "        # Eigendecomposition\n",
+    "        eigenvals, _ = torch.linalg.eigh(gram)\n",
+    "        eigen_mean = eigenvals.mean(dim=1)\n",
+    "        eigen_std = eigenvals.std(dim=1)\n",
+    "        eigen_max = eigenvals.max(dim=1)[0]\n",
+    "        eigen_min = torch.clamp(eigenvals.min(dim=1)[0], min=1e-8)\n",
+    "        eigen_cond = eigen_max / eigen_min\n",
+    "        eigen_trace = eigenvals.sum(dim=1)\n",
+    "        eigen_norm = torch.linalg.norm(eigenvals, dim=1)\n",
+    "        eigen_entropy = -(eigenvals / (eigenvals.sum(dim=1, keepdim=True) + 1e-8) *\n",
+    "                          torch.log(eigenvals / (eigenvals.sum(dim=1, keepdim=True) + 1e-8) + 1e-8)).sum(dim=1)\n",
+    "        eigen_rank = (eigenvals > 1e-5).sum(dim=1).float()\n",
+    "        eigen_sparsity = (eigenvals < 1e-3).sum(dim=1).float() / self.gram_dim\n",
+    "\n",
+    "        # SVD\n",
+    "        U, S, Vh = torch.svd(gram)\n",
+    "        svd_mean = S.mean(dim=1)\n",
+    "        svd_std = S.std(dim=1)\n",
+    "        svd_ratio = S[:, 0] / (S[:, -1] + 1e-8)\n",
+    "        svd_sum = S.sum(dim=1)\n",
+    "        svd_norm = torch.linalg.norm(S, dim=1)\n",
+    "        svd_rank = (S > 1e-5).sum(dim=1).float()\n",
+    "        svd_cumsum = (S.cumsum(dim=1) / S.sum(dim=1, keepdim=True))[:, -1]\n",
+    "\n",
+    "        # Cholesky\n",
+    "        try:\n",
+    "            L = torch.linalg.cholesky(gram)\n",
+    "            chol_diag = torch.diagonal(L, dim1=1, dim2=2)\n",
+    "            chol_mean = chol_diag.mean(dim=1)\n",
+    "            chol_std = chol_diag.std(dim=1)\n",
+    "            chol_sum = chol_diag.sum(dim=1)\n",
+    "            chol_norm = torch.linalg.norm(chol_diag, dim=1)\n",
+    "            chol_prod = torch.prod(chol_diag, dim=1)\n",
+    "            chol_trace = L.diagonal(dim1=1, dim2=2).sum(dim=1)\n",
+    "        except:\n",
+    "            # Fallback if Cholesky fails\n",
+    "            chol_mean = torch.zeros(B, device=gram.device)\n",
+    "            chol_std = torch.zeros(B, device=gram.device)\n",
+    "            chol_sum = torch.zeros(B, device=gram.device)\n",
+    "            chol_norm = torch.zeros(B, device=gram.device)\n",
+    "            chol_prod = torch.ones(B, device=gram.device)\n",
+    "            chol_trace = torch.zeros(B, device=gram.device)\n",
+    "\n",
+    "        # Combine all signatures\n",
+    "        eigen_sig = torch.stack([eigen_mean, eigen_std, eigen_max, eigen_min, eigen_cond,\n",
+    "                                 eigen_trace, eigen_norm, eigen_entropy, eigen_rank, eigen_sparsity], dim=1)\n",
+    "        svd_sig = torch.stack([svd_mean, svd_std, svd_ratio, svd_sum, svd_norm,\n",
+    "                               svd_rank, svd_cumsum,\n",
+    "                               S[:, 0], S[:, 1] if S.shape[1] > 1 else S[:, 0],\n",
+    "                               (S[:, :5].sum(dim=1) / S.sum(dim=1))], dim=1)\n",
+    "        chol_sig = torch.stack([chol_mean, chol_std, chol_sum, chol_norm, chol_prod,\n",
+    "                                chol_trace, chol_prod.log(),\n",
+    "                                chol_diag.max(dim=1)[0], chol_diag.min(dim=1)[0],\n",
+    "                                chol_diag.max(dim=1)[0] / (chol_diag.min(dim=1)[0] + 1e-8)], dim=1)\n",
+    "\n",
+    "        sig_combined = torch.cat([eigen_sig, svd_sig, chol_sig], dim=1)  # (B, 30)\n",
+    "        combined = torch.cat([feat_red, sig_combined], dim=1)\n",
+    "\n",
+    "        emb = l2_normalize(self.embed(combined))\n",
+    "        logits, tri = self.classifier(emb)\n",
+    "        return logits, emb\n",
+    "\n",
+    "model = EigenSignatureModel(embed_dim=128, num_classes=10).to(device)\n",
+    "print(f\"[5.2c] Model params: {sum(p.numel() for p in model.parameters()):,}\")\n",
+    "\n",
+    "result_5_2c = run_experiment(\n",
+    "    \"5.2c_Eigen_SVD_Cholesky\",\n",
+    "    model, train_loader, val_loader, device,\n",
+    "    epochs=100, lr=1e-3, weight_decay=1e-4, patience=20\n",
+    ")\n",
+    "scoreboard.add(result_5_2c)\n",
+    ""
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 5.2d: NMF of Magnitude Spectrum\n",
+    "\n",
+    "Parts-based decomposition using sklearn NMF"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "\n",
+    "from sklearn.decomposition import NMF\n",
+    "print(\"[5.2d] NMF of Magnitude Spectrum\")\n",
+    "print(\"=\"*70)\n",
+    "\n",
+    "class NMFDecompositionModel(nn.Module):\n",
+    "    '''STFT magnitude \u2192 NMF parts-based decomposition.'''\n",
+    "    def __init__(self, embed_dim=128, num_classes=10, n_components=16):\n",
+    "        super().__init__()\n",
+    "        self.stft = STFTFrontend()\n",
+    "        self.n_components = n_components\n",
+    "        self.nmf_model = NMF(n_components=n_components, init='random', random_state=42, max_iter=200)\n",
+    "\n",
+    "        self.feat_dim = n_components + 10  # NMF coefficients + stats\n",
+    "        self.embed = nn.Linear(self.feat_dim, embed_dim)\n",
+    "        self.classifier = ConstellationClassifier(embed_dim, num_anchors=64, num_classes=num_classes)\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        stft_feat = self.stft(x)  # (B, 1536)\n",
+    "\n",
+    "        # NMF on magnitude (treat as 1536-dim features)\n",
+    "        stft_np = stft_feat.detach().cpu().numpy()\n",
+    "        stft_np = np.abs(stft_np)  # Ensure non-negative\n",
+    "\n",
+    "        nmf_coeff = torch.zeros(x.shape[0], self.n_components, device=x.device)\n",
+    "        try:\n",
+    "            for b in range(x.shape[0]):\n",
+    "                coeff = self.nmf_model.fit_transform(stft_np[b:b+1])\n",
+    "                nmf_coeff[b] = torch.from_numpy(coeff[0]).float().to(x.device)\n",
+    "        except:\n",
+    "            # If NMF fails, use zeros\n",
+    "            pass\n",
+    "\n",
+    "        # Extract statistics\n",
+    "        nmf_mean = nmf_coeff.mean(dim=1, keepdim=True)\n",
+    "        nmf_std = nmf_coeff.std(dim=1, keepdim=True)\n",
+    "        nmf_max = nmf_coeff.max(dim=1)[0].unsqueeze(1)\n",
+    "        nmf_norm = torch.linalg.norm(nmf_coeff, dim=1, keepdim=True)\n",
+    "        sparsity = (nmf_coeff < 0.01).sum(dim=1, keepdim=True).float() / self.n_components\n",
+    "        energy = (nmf_coeff ** 2).sum(dim=1, keepdim=True)\n",
+    "        entropy = -(nmf_coeff / (nmf_coeff.sum(dim=1, keepdim=True) + 1e-8) *\n",
+    "                    torch.log(nmf_coeff / (nmf_coeff.sum(dim=1, keepdim=True) + 1e-8) + 1e-8)).sum(dim=1, keepdim=True)\n",
+    "        active = (nmf_coeff > 1e-3).sum(dim=1, keepdim=True).float()\n",
+    "\n",
+    "        stats = torch.cat([nmf_mean, nmf_std, nmf_max, nmf_norm, sparsity, energy, entropy, active], dim=1)  # (B, 8)\n",
+    "        combined = torch.cat([nmf_coeff, stats], dim=1)  # (B, 24)\n",
+    "\n",
+    "        emb = l2_normalize(self.embed(combined))\n",
+    "        logits, tri = self.classifier(emb)\n",
+    "        return logits, emb\n",
+    "\n",
+    "model = NMFDecompositionModel(embed_dim=128, num_classes=10, n_components=16).to(device)\n",
+    "print(f\"[5.2d] Model params: {sum(p.numel() for p in model.parameters()):,}\")\n",
+    "\n",
+    "result_5_2d = run_experiment(\n",
+    "    \"5.2d_NMF_Decomposition\",\n",
+    "    model, train_loader, val_loader, device,\n",
+    "    epochs=100, lr=1e-3, weight_decay=1e-4, patience=20\n",
+    ")\n",
+    "scoreboard.add(result_5_2d)\n",
+    ""
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 5.2e: Tucker Tensor Decomposition\n",
+    "\n",
+    "3D structure (H, W, freq_bins) per color channel, core tensor encodes interactions"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "\n",
+    "print(\"[5.2e] Tucker Tensor Decomposition\")\n",
+    "print(\"=\"*70)\n",
+    "\n",
+    "class TuckerDecompositionModel(nn.Module):\n",
+    "    '''Reshape STFT features to 3D tensor, apply Tucker decomposition.'''\n",
+    "    def __init__(self, embed_dim=128, num_classes=10):\n",
+    "        super().__init__()\n",
+    "        # Manually implement Tucker-like feature extraction (full Tucker is complex)\n",
+    "        self.stft = STFTFrontend()\n",
+    "        self.dim_reduce = nn.Linear(1536, 256)\n",
+    "\n",
+    "        self.feat_dim = 256 + 20  # Reduced features + Tucker stats\n",
+    "        self.embed = nn.Linear(self.feat_dim, embed_dim)\n",
+    "        self.classifier = ConstellationClassifier(embed_dim, num_anchors=64, num_classes=num_classes)\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        stft_feat = self.stft(x)  # (B, 1536)\n",
+    "        feat_red = self.dim_reduce(stft_feat)  # (B, 256)\n",
+    "\n",
+    "        B = feat_red.shape[0]\n",
+    "\n",
+    "        # Reshape to 3D: (B, 8, 32, 32) approximation\n",
+    "        tensor_3d = feat_red.view(B, 8, 32)  # Approximate 3D\n",
+    "\n",
+    "        # Tucker-like: compute mode-wise statistics\n",
+    "        mode0_unfold = tensor_3d.view(B, 8, -1)  # (B, 8, 32)\n",
+    "        mode1_unfold = tensor_3d.transpose(1, 2)  # (B, 32, 8)\n",
+    "\n",
+    "        # Covariances along each mode\n",
+    "        cov0 = torch.bmm(mode0_unfold, mode0_unfold.transpose(1, 2))  # (B, 8, 8)\n",
+    "        cov1 = torch.bmm(mode1_unfold, mode1_unfold.transpose(1, 2))  # (B, 32, 32)\n",
+    "\n",
+    "        # Extract eigenvalues from mode covariances\n",
+    "        eig0, _ = torch.linalg.eigh(cov0)  # (B, 8)\n",
+    "        eig1, _ = torch.linalg.eigh(cov1)  # (B, 32)\n",
+    "\n",
+    "        # Core tensor stats\n",
+    "        core_mean = eig0.mean(dim=1)\n",
+    "        core_std = eig0.std(dim=1)\n",
+    "        core_max = eig0.max(dim=1)[0]\n",
+    "        mode_mean = eig1.mean(dim=1)\n",
+    "        mode_std = eig1.std(dim=1)\n",
+    "        mode_trace = eig1.sum(dim=1)\n",
+    "\n",
+    "        # Mode interaction (trace of Gram across modes)\n",
+    "        interaction = (eig0 * eig0.mean(dim=1, keepdim=True)).sum(dim=1)\n",
+    "        contraction = torch.einsum('bi,bi->b', eig0.view(B, -1), eig1[:, :8])\n",
+    "\n",
+    "        # Coupled features\n",
+    "        coupling1 = (eig0[:, 0] + eig0[:, -1]) / 2\n",
+    "        coupling2 = (eig1[:, 0] + eig1[:, -1]) / 2\n",
+    "        coupling3 = eig0.max(dim=1)[0] / (eig0.min(dim=1)[0] + 1e-8)\n",
+    "        coupling4 = eig1.max(dim=1)[0] / (eig1.min(dim=1)[0] + 1e-8)\n",
+    "\n",
+    "        tucker_stats = torch.stack([core_mean, core_std, core_max, mode_mean, mode_std,\n",
+    "                                     mode_trace, interaction, contraction,\n",
+    "                                     coupling1, coupling2, coupling3, coupling4,\n",
+    "                                     eig0.sum(dim=1), eig1.sum(dim=1),\n",
+    "                                     torch.linalg.norm(eig0, dim=1),\n",
+    "                                     torch.linalg.norm(eig1, dim=1),\n",
+    "                                     (eig0 > 1e-5).sum(dim=1).float(),\n",
+    "                                     (eig1 > 1e-5).sum(dim=1).float(),\n",
+    "                                     eig0.max(dim=1)[0] - eig0.min(dim=1)[0],\n",
+    "                                     eig1.max(dim=1)[0] - eig1.min(dim=1)[0]], dim=1)  # (B, 20)\n",
+    "\n",
+    "        combined = torch.cat([feat_red, tucker_stats], dim=1)\n",
+    "\n",
+    "        emb = l2_normalize(self.embed(combined))\n",
+    "        logits, tri = self.classifier(emb)\n",
+    "        return logits, emb\n",
+    "\n",
+    "model = TuckerDecompositionModel(embed_dim=128, num_classes=10).to(device)\n",
+    "print(f\"[5.2e] Model params: {sum(p.numel() for p in model.parameters()):,}\")\n",
+    "\n",
+    "result_5_2e = run_experiment(\n",
+    "    \"5.2e_Tucker_Decomposition\",\n",
+    "    model, train_loader, val_loader, device,\n",
+    "    epochs=100, lr=1e-3, weight_decay=1e-4, patience=20\n",
+    ")\n",
+    "scoreboard.add(result_5_2e)\n",
+    ""
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Final Scoreboard"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "\n",
+    "print(\"\\n\" + \"=\"*80)\n",
+    "print(\"EXPERIMENT 5 COMPLETE: Matrix Decomposition Signatures\")\n",
+    "print(\"=\"*80)\n",
+    "scoreboard._print_table()\n",
+    "scoreboard.save(\"scoreboard_exp5.json\")\n",
+    ""
+   ],
+   "execution_count": null,
+   "outputs": []
+  }
+ ]
+}

spectral/experiment_6_losses_and_anchors.ipynb

@@ -0,0 +1,1364 @@
+{
+ "nbformat": 4,
+ "nbformat_minor": 5,
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.10.0"
+  },
+  "accelerator": "GPU",
+  "colab": {
+   "provenance": [],
+   "gpuType": "T4"
+  }
+ },
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Experiment 6: Information-Theoretic Losses & Anchor Management\n",
+    "\n",
+    "Advanced loss functions and anchor manipulation strategies."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "\n",
+    "!pip install -q kymatio scipy tensorboard scikit-learn matplotlib numpy torch torchvision torchaudio scipy\n",
+    ""
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Shared Utilities"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "\n",
+    "# \u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\n",
+    "# SHARED UTILITIES \u2014 GeoLIP Hypersphere Experiments\n",
+    "# \u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\n",
+    "import os, sys, time, math, json, datetime, warnings\n",
+    "import numpy as np\n",
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "import torch.nn.functional as F\n",
+    "from torch.utils.data import DataLoader\n",
+    "from torch.utils.tensorboard import SummaryWriter\n",
+    "import torchvision\n",
+    "import torchvision.transforms as T\n",
+    "import matplotlib\n",
+    "matplotlib.use(\"Agg\")\n",
+    "import matplotlib.pyplot as plt\n",
+    "from collections import defaultdict\n",
+    "\n",
+    "warnings.filterwarnings(\"ignore\", category=UserWarning)\n",
+    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
+    "print(f\"[DEVICE] {device}\" + (f\" \u2014 {torch.cuda.get_device_name()}\" if device.type == \"cuda\" else \"\"))\n",
+    "\n",
+    "# \u2500\u2500 CIFAR-10 Loaders \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
+    "CIFAR10_MEAN = (0.4914, 0.4822, 0.4465)\n",
+    "CIFAR10_STD  = (0.2470, 0.2435, 0.2616)\n",
+    "CIFAR10_CLASSES = (\"plane\",\"car\",\"bird\",\"cat\",\"deer\",\"dog\",\"frog\",\"horse\",\"ship\",\"truck\")\n",
+    "\n",
+    "def get_cifar10_loaders(batch_size=256, num_workers=2, augment=True):\n",
+    "    \"\"\"Return train_loader, val_loader for CIFAR-10.\"\"\"\n",
+    "    tf_train = T.Compose([\n",
+    "        T.RandomCrop(32, padding=4),\n",
+    "        T.RandomHorizontalFlip(),\n",
+    "        T.ToTensor(),\n",
+    "        T.Normalize(CIFAR10_MEAN, CIFAR10_STD),\n",
+    "    ]) if augment else T.Compose([T.ToTensor(), T.Normalize(CIFAR10_MEAN, CIFAR10_STD)])\n",
+    "    tf_val = T.Compose([T.ToTensor(), T.Normalize(CIFAR10_MEAN, CIFAR10_STD)])\n",
+    "\n",
+    "    train_ds = torchvision.datasets.CIFAR10(root=\"./data\", train=True,  download=True, transform=tf_train)\n",
+    "    val_ds   = torchvision.datasets.CIFAR10(root=\"./data\", train=False, download=True, transform=tf_val)\n",
+    "    train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True,  num_workers=num_workers, pin_memory=True, drop_last=True)\n",
+    "    val_loader   = DataLoader(val_ds,   batch_size=batch_size, shuffle=False, num_workers=num_workers, pin_memory=True)\n",
+    "    print(f\"[DATA] CIFAR-10: {len(train_ds)} train, {len(val_ds)} val, batch_size={batch_size}\")\n",
+    "    return train_loader, val_loader\n",
+    "\n",
+    "def get_cifar10_raw_loaders(batch_size=256, num_workers=2):\n",
+    "    \"\"\"Return loaders with only ToTensor (no normalization) for custom front-ends.\"\"\"\n",
+    "    tf = T.Compose([T.ToTensor()])\n",
+    "    train_ds = torchvision.datasets.CIFAR10(root=\"./data\", train=True,  download=True, transform=tf)\n",
+    "    val_ds   = torchvision.datasets.CIFAR10(root=\"./data\", train=False, download=True, transform=tf)\n",
+    "    train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True,  num_workers=num_workers, pin_memory=True, drop_last=True)\n",
+    "    val_loader   = DataLoader(val_ds,   batch_size=batch_size, shuffle=False, num_workers=num_workers, pin_memory=True)\n",
+    "    print(f\"[DATA] CIFAR-10 (raw): {len(train_ds)} train, {len(val_ds)} val\")\n",
+    "    return train_loader, val_loader\n",
+    "\n",
+    "# \u2500\u2500 L2 Normalization to Hypersphere \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
+    "def l2_normalize(x, dim=-1, eps=1e-8):\n",
+    "    \"\"\"Project to unit hypersphere S^{d-1}.\"\"\"\n",
+    "    return F.normalize(x, p=2, dim=dim, eps=eps)\n",
+    "\n",
+    "# \u2500\u2500 Constellation Module \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
+    "class Constellation(nn.Module):\n",
+    "    \"\"\"\n",
+    "    K anchors on S^{d-1}. Triangulation: compute distance from input to each anchor.\n",
+    "    Output: K-dim vector of (1 - cos_sim) distances = triangulation profile.\n",
+    "    \"\"\"\n",
+    "    def __init__(self, embed_dim, num_anchors=64, init=\"uniform\"):\n",
+    "        super().__init__()\n",
+    "        self.embed_dim = embed_dim\n",
+    "        self.num_anchors = num_anchors\n",
+    "        anchors = self._init_anchors(embed_dim, num_anchors, init)\n",
+    "        self.anchors = nn.Parameter(anchors, requires_grad=False)  # anchors are NOT learned\n",
+    "\n",
+    "    def _init_anchors(self, d, K, method):\n",
+    "        if method == \"uniform\":\n",
+    "            # QR-based uniform hypersphere initialization\n",
+    "            A = torch.randn(K, d)\n",
+    "            Q, _ = torch.linalg.qr(A.T)\n",
+    "            anchors = Q.T[:K]\n",
+    "            # If K > d, fill remaining with random unit vectors\n",
+    "            if K > d:\n",
+    "                extra = F.normalize(torch.randn(K - d, d), dim=-1)\n",
+    "                anchors = torch.cat([anchors, extra], dim=0)\n",
+    "            return F.normalize(anchors, dim=-1)\n",
+    "        elif method == \"random\":\n",
+    "            return F.normalize(torch.randn(K, d), dim=-1)\n",
+    "        else:\n",
+    "            return F.normalize(torch.randn(K, d), dim=-1)\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        \"\"\"x: (B, d) on S^{d-1} \u2192 (B, K) triangulation distances.\"\"\"\n",
+    "        x_norm = l2_normalize(x)\n",
+    "        cos_sim = torch.mm(x_norm, self.anchors.T)  # (B, K)\n",
+    "        distances = 1.0 - cos_sim  # in [0, 2]\n",
+    "        return distances\n",
+    "\n",
+    "    def push_anchors(self, embeddings, labels, num_classes=10, momentum=0.1):\n",
+    "        \"\"\"Push anchors toward class centroids (call periodically).\"\"\"\n",
+    "        with torch.no_grad():\n",
+    "            for c in range(num_classes):\n",
+    "                mask = (labels == c)\n",
+    "                if mask.sum() < 2:\n",
+    "                    continue\n",
+    "                centroid = l2_normalize(embeddings[mask].mean(dim=0, keepdim=True))\n",
+    "                # Assign anchors to this class based on nearest centroid\n",
+    "                start = c * (self.num_anchors // num_classes)\n",
+    "                end = start + (self.num_anchors // num_classes)\n",
+    "                self.anchors.data[start:end] = l2_normalize(\n",
+    "                    (1 - momentum) * self.anchors.data[start:end] + momentum * centroid\n",
+    "                )\n",
+    "\n",
+    "    def anchor_activation(self, x):\n",
+    "        \"\"\"Return fraction of anchors that are 'active' (nearest to at least one sample).\"\"\"\n",
+    "        x_norm = l2_normalize(x)\n",
+    "        cos_sim = torch.mm(x_norm, self.anchors.T)\n",
+    "        nearest = cos_sim.argmax(dim=1)\n",
+    "        active = len(nearest.unique())\n",
+    "        return active, self.num_anchors\n",
+    "\n",
+    "# \u2500\u2500 Patchwork Block \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
+    "class SquaredReLU(nn.Module):\n",
+    "    def forward(self, x):\n",
+    "        return F.relu(x) ** 2\n",
+    "\n",
+    "class PatchworkBlock(nn.Module):\n",
+    "    \"\"\"SquaredReLU + LayerNorm block for post-constellation processing.\"\"\"\n",
+    "    def __init__(self, dim):\n",
+    "        super().__init__()\n",
+    "        self.net = nn.Sequential(\n",
+    "            nn.Linear(dim, dim),\n",
+    "            SquaredReLU(),\n",
+    "            nn.LayerNorm(dim),\n",
+    "            nn.Linear(dim, dim),\n",
+    "            SquaredReLU(),\n",
+    "            nn.LayerNorm(dim),\n",
+    "        )\n",
+    "    def forward(self, x):\n",
+    "        return x + self.net(x)  # residual\n",
+    "\n",
+    "class Patchwork(nn.Module):\n",
+    "    \"\"\"Stack of PatchworkBlocks.\"\"\"\n",
+    "    def __init__(self, dim, num_blocks=2):\n",
+    "        super().__init__()\n",
+    "        self.blocks = nn.Sequential(*[PatchworkBlock(dim) for _ in range(num_blocks)])\n",
+    "    def forward(self, x):\n",
+    "        return self.blocks(x)\n",
+    "\n",
+    "# \u2500\u2500 Classifier Head \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
+    "class ConstellationClassifier(nn.Module):\n",
+    "    \"\"\"Full pipeline: encoder output \u2192 constellation \u2192 patchwork \u2192 classify.\"\"\"\n",
+    "    def __init__(self, embed_dim, num_anchors=64, num_classes=10, patchwork_blocks=2):\n",
+    "        super().__init__()\n",
+    "        self.constellation = Constellation(embed_dim, num_anchors)\n",
+    "        self.patchwork = Patchwork(num_anchors, patchwork_blocks)\n",
+    "        self.head = nn.Linear(num_anchors, num_classes)\n",
+    "\n",
+    "    def forward(self, embeddings):\n",
+    "        tri = self.constellation(embeddings)       # (B, K)\n",
+    "        pw  = self.patchwork(tri)                   # (B, K)\n",
+    "        logits = self.head(pw)                      # (B, 10)\n",
+    "        return logits, tri\n",
+    "\n",
+    "    def count_params(self):\n",
+    "        return sum(p.numel() for p in self.parameters() if p.requires_grad)\n",
+    "\n",
+    "# \u2500\u2500 Losses \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
+    "class CayleyMengerVolumeLoss(nn.Module):\n",
+    "    \"\"\"CV loss: encourage constellation volume in 0.18-0.25 band.\"\"\"\n",
+    "    def __init__(self, target_low=0.18, target_high=0.25):\n",
+    "        super().__init__()\n",
+    "        self.target_low = target_low\n",
+    "        self.target_high = target_high\n",
+    "\n",
+    "    def forward(self, tri):\n",
+    "        \"\"\"tri: (B, K) triangulation distances.\"\"\"\n",
+    "        # Approximate volume via determinant of distance submatrix\n",
+    "        B, K = tri.shape\n",
+    "        # Use std of triangulation as proxy for volume spread\n",
+    "        vol = tri.std(dim=1).mean()\n",
+    "        # Penalty if outside target band\n",
+    "        below = F.relu(self.target_low - vol)\n",
+    "        above = F.relu(vol - self.target_high)\n",
+    "        return below + above, vol.item()\n",
+    "\n",
+    "class InfoNCELoss(nn.Module):\n",
+    "    \"\"\"Self-supervised contrastive loss on augmented views.\"\"\"\n",
+    "    def __init__(self, temperature=0.07):\n",
+    "        super().__init__()\n",
+    "        self.temperature = temperature\n",
+    "\n",
+    "    def forward(self, z1, z2):\n",
+    "        \"\"\"z1, z2: (B, d) L2-normalized embeddings of two augmented views.\"\"\"\n",
+    "        z1 = l2_normalize(z1)\n",
+    "        z2 = l2_normalize(z2)\n",
+    "        B = z1.shape[0]\n",
+    "        z = torch.cat([z1, z2], dim=0)  # (2B, d)\n",
+    "        sim = torch.mm(z, z.T) / self.temperature  # (2B, 2B)\n",
+    "        # Mask out self-similarity\n",
+    "        mask = ~torch.eye(2*B, dtype=bool, device=z.device)\n",
+    "        sim = sim.masked_fill(~mask, -1e9)\n",
+    "        # Positive pairs: (i, i+B) and (i+B, i)\n",
+    "        labels = torch.cat([torch.arange(B, 2*B), torch.arange(B)]).to(z.device)\n",
+    "        loss = F.cross_entropy(sim, labels)\n",
+    "        return loss\n",
+    "\n",
+    "class ClassConditionalInfoNCE(nn.Module):\n",
+    "    \"\"\"Class-conditional contrastive: positive = same class, not same image.\"\"\"\n",
+    "    def __init__(self, temperature=0.07):\n",
+    "        super().__init__()\n",
+    "        self.temperature = temperature\n",
+    "\n",
+    "    def forward(self, z, labels):\n",
+    "        z = l2_normalize(z)\n",
+    "        B = z.shape[0]\n",
+    "        sim = torch.mm(z, z.T) / self.temperature\n",
+    "        # Mask: positive if same class, negative otherwise\n",
+    "        label_mask = (labels.unsqueeze(0) == labels.unsqueeze(1)).float()\n",
+    "        label_mask.fill_diagonal_(0)  # exclude self\n",
+    "        # For each sample, compute log-softmax over all non-self, weighted by positives\n",
+    "        logits = sim - torch.diag(torch.full((B,), 1e9, device=z.device))\n",
+    "        pos_mask = label_mask\n",
+    "        neg_mask = 1 - label_mask\n",
+    "        neg_mask.fill_diagonal_(0)\n",
+    "        # InfoNCE: -log(sum_pos exp(sim) / sum_all exp(sim))\n",
+    "        exp_sim = torch.exp(logits)\n",
+    "        pos_sum = (exp_sim * pos_mask).sum(dim=1)\n",
+    "        all_sum = (exp_sim * (pos_mask + neg_mask)).sum(dim=1)\n",
+    "        loss = -torch.log(pos_sum / (all_sum + 1e-8) + 1e-8).mean()\n",
+    "        return loss\n",
+    "\n",
+    "# \u2500\u2500 TensorBoard Helpers \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
+    "def make_writer(experiment_id, base_dir=\"runs\"):\n",
+    "    \"\"\"Create a TensorBoard SummaryWriter with timestamped run directory.\"\"\"\n",
+    "    timestamp = datetime.datetime.now().strftime(\"%Y%m%d_%H%M%S\")\n",
+    "    log_dir = os.path.join(base_dir, f\"{experiment_id}_{timestamp}\")\n",
+    "    writer = SummaryWriter(log_dir=log_dir)\n",
+    "    print(f\"[TB] Logging to {log_dir}\")\n",
+    "    return writer\n",
+    "\n",
+    "def log_epoch_metrics(writer, epoch, metrics_dict, prefix=\"\"):\n",
+    "    \"\"\"Log a dict of metrics to TensorBoard.\"\"\"\n",
+    "    for key, val in metrics_dict.items():\n",
+    "        tag = f\"{prefix}/{key}\" if prefix else key\n",
+    "        writer.add_scalar(tag, val, epoch)\n",
+    "\n",
+    "def log_embedding_viz(writer, embeddings, labels, epoch, tag=\"embeddings\"):\n",
+    "    \"\"\"Log embeddings for TensorBoard projector.\"\"\"\n",
+    "    meta = [CIFAR10_CLASSES[l] for l in labels]\n",
+    "    writer.add_embedding(embeddings, metadata=meta, global_step=epoch, tag=tag)\n",
+    "\n",
+    "def log_anchor_histogram(writer, constellation, epoch):\n",
+    "    \"\"\"Log anchor norm distribution.\"\"\"\n",
+    "    norms = constellation.anchors.data.norm(dim=1)\n",
+    "    writer.add_histogram(\"anchors/norms\", norms, epoch)\n",
+    "    cos_sim_matrix = torch.mm(constellation.anchors.data, constellation.anchors.data.T)\n",
+    "    off_diag = cos_sim_matrix[~torch.eye(cos_sim_matrix.shape[0], dtype=bool)].flatten()\n",
+    "    writer.add_histogram(\"anchors/inter_cosine_sim\", off_diag, epoch)\n",
+    "\n",
+    "def log_confusion_matrix(writer, all_preds, all_labels, epoch):\n",
+    "    \"\"\"Log confusion matrix as image to TensorBoard.\"\"\"\n",
+    "    from sklearn.metrics import confusion_matrix\n",
+    "    cm = confusion_matrix(all_labels, all_preds, labels=range(10))\n",
+    "    fig, ax = plt.subplots(1, 1, figsize=(8, 8))\n",
+    "    im = ax.imshow(cm, interpolation=\"nearest\", cmap=plt.cm.Blues)\n",
+    "    ax.set_title(f\"Confusion Matrix \u2014 Epoch {epoch}\")\n",
+    "    ax.set_xlabel(\"Predicted\")\n",
+    "    ax.set_ylabel(\"True\")\n",
+    "    tick_marks = range(10)\n",
+    "    ax.set_xticks(tick_marks)\n",
+    "    ax.set_xticklabels(CIFAR10_CLASSES, rotation=45, ha=\"right\")\n",
+    "    ax.set_yticks(tick_marks)\n",
+    "    ax.set_yticklabels(CIFAR10_CLASSES)\n",
+    "    for i in range(10):\n",
+    "        for j in range(10):\n",
+    "            ax.text(j, i, str(cm[i, j]), ha=\"center\", va=\"center\",\n",
+    "                    color=\"white\" if cm[i, j] > cm.max() / 2 else \"black\", fontsize=7)\n",
+    "    fig.tight_layout()\n",
+    "    writer.add_figure(\"confusion_matrix\", fig, epoch)\n",
+    "    plt.close(fig)\n",
+    "\n",
+    "def log_triangulation_stats(writer, tri, epoch):\n",
+    "    \"\"\"Log triangulation distance statistics.\"\"\"\n",
+    "    writer.add_scalar(\"tri/mean\", tri.mean().item(), epoch)\n",
+    "    writer.add_scalar(\"tri/std\", tri.std().item(), epoch)\n",
+    "    writer.add_scalar(\"tri/min\", tri.min().item(), epoch)\n",
+    "    writer.add_scalar(\"tri/max\", tri.max().item(), epoch)\n",
+    "    writer.add_histogram(\"tri/distances\", tri.detach().cpu(), epoch)\n",
+    "\n",
+    "def log_scatter_2d(writer, embeddings, labels, epoch, tag=\"scatter_2d\", method=\"pca\"):\n",
+    "    \"\"\"Log 2D scatter plot of embeddings projected via PCA or t-SNE.\"\"\"\n",
+    "    emb_np = embeddings.detach().cpu().numpy()\n",
+    "    lab_np = labels.detach().cpu().numpy() if torch.is_tensor(labels) else np.array(labels)\n",
+    "    if method == \"pca\":\n",
+    "        from sklearn.decomposition import PCA\n",
+    "        proj = PCA(n_components=2).fit_transform(emb_np)\n",
+    "    else:\n",
+    "        from sklearn.manifold import TSNE\n",
+    "        proj = TSNE(n_components=2, perplexity=30, n_iter=500).fit_transform(emb_np)\n",
+    "    fig, ax = plt.subplots(1, 1, figsize=(8, 8))\n",
+    "    for c in range(10):\n",
+    "        mask = lab_np == c\n",
+    "        ax.scatter(proj[mask, 0], proj[mask, 1], s=3, alpha=0.5, label=CIFAR10_CLASSES[c])\n",
+    "    ax.legend(markerscale=5, fontsize=8)\n",
+    "    ax.set_title(f\"{tag} ({method.upper()}) \u2014 Epoch {epoch}\")\n",
+    "    fig.tight_layout()\n",
+    "    writer.add_figure(f\"{tag}/{method}\", fig, epoch)\n",
+    "    plt.close(fig)\n",
+    "\n",
+    "# \u2500\u2500 Training Utilities \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
+    "def train_one_epoch(model, loader, optimizer, criterion, device, epoch=0,\n",
+    "                    scheduler=None, extra_losses=None):\n",
+    "    \"\"\"\n",
+    "    Generic training loop for one epoch.\n",
+    "    model: callable(batch) \u2192 (logits, embeddings_on_sphere)\n",
+    "    extra_losses: list of (name, loss_fn, weight) where loss_fn(embeddings) \u2192 scalar\n",
+    "    \"\"\"\n",
+    "    model.train()\n",
+    "    total_loss = 0\n",
+    "    total_ce = 0\n",
+    "    correct = 0\n",
+    "    total = 0\n",
+    "    loss_components = defaultdict(float)\n",
+    "\n",
+    "    for batch_idx, (images, labels) in enumerate(loader):\n",
+    "        images, labels = images.to(device), labels.to(device)\n",
+    "        optimizer.zero_grad()\n",
+    "\n",
+    "        logits, embeddings = model(images)\n",
+    "        ce_loss = criterion(logits, labels)\n",
+    "        loss = ce_loss\n",
+    "        loss_components[\"CE\"] += ce_loss.item()\n",
+    "\n",
+    "        if extra_losses:\n",
+    "            for name, loss_fn, weight in extra_losses:\n",
+    "                if name == \"CV\":\n",
+    "                    el, vol_val = loss_fn(embeddings)\n",
+    "                    loss_components[\"CV_vol\"] += vol_val\n",
+    "                elif name == \"InfoNCE_class\":\n",
+    "                    el = loss_fn(embeddings, labels)\n",
+    "                else:\n",
+    "                    el = loss_fn(embeddings)\n",
+    "                loss = loss + weight * el\n",
+    "                loss_components[name] += el.item()\n",
+    "\n",
+    "        loss.backward()\n",
+    "        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)\n",
+    "        optimizer.step()\n",
+    "        if scheduler:\n",
+    "            scheduler.step()\n",
+    "\n",
+    "        total_loss += loss.item()\n",
+    "        _, predicted = logits.max(1)\n",
+    "        total += labels.size(0)\n",
+    "        correct += predicted.eq(labels).sum().item()\n",
+    "\n",
+    "    n_batches = len(loader)\n",
+    "    metrics = {\n",
+    "        \"loss/total\": total_loss / n_batches,\n",
+    "        \"loss/CE\": loss_components[\"CE\"] / n_batches,\n",
+    "        \"accuracy/train\": 100.0 * correct / total,\n",
+    "    }\n",
+    "    for k, v in loss_components.items():\n",
+    "        if k not in (\"CE\",):\n",
+    "            metrics[f\"loss/{k}\"] = v / n_batches\n",
+    "    return metrics\n",
+    "\n",
+    "@torch.no_grad()\n",
+    "def evaluate(model, loader, criterion, device, collect_embeddings=False):\n",
+    "    \"\"\"Evaluate model on validation set. Optionally collect embeddings.\"\"\"\n",
+    "    model.eval()\n",
+    "    total_loss = 0\n",
+    "    correct = 0\n",
+    "    total = 0\n",
+    "    all_preds = []\n",
+    "    all_labels = []\n",
+    "    all_embeddings = []\n",
+    "\n",
+    "    for images, labels in loader:\n",
+    "        images, labels = images.to(device), labels.to(device)\n",
+    "        logits, embeddings = model(images)\n",
+    "        loss = criterion(logits, labels)\n",
+    "        total_loss += loss.item()\n",
+    "        _, predicted = logits.max(1)\n",
+    "        total += labels.size(0)\n",
+    "        correct += predicted.eq(labels).sum().item()\n",
+    "        all_preds.extend(predicted.cpu().numpy())\n",
+    "        all_labels.extend(labels.cpu().numpy())\n",
+    "        if collect_embeddings:\n",
+    "            all_embeddings.append(embeddings.cpu())\n",
+    "\n",
+    "    n_batches = len(loader)\n",
+    "    metrics = {\n",
+    "        \"loss/val\": total_loss / n_batches,\n",
+    "        \"accuracy/val\": 100.0 * correct / total,\n",
+    "    }\n",
+    "    result = {\"metrics\": metrics, \"preds\": np.array(all_preds), \"labels\": np.array(all_labels)}\n",
+    "    if collect_embeddings:\n",
+    "        result[\"embeddings\"] = torch.cat(all_embeddings, dim=0)\n",
+    "    return result\n",
+    "\n",
+    "def run_experiment(experiment_id, model, train_loader, val_loader, device,\n",
+    "                   epochs=100, lr=1e-3, weight_decay=1e-4, extra_losses=None,\n",
+    "                   push_anchors_every=50, log_embeddings_every=25,\n",
+    "                   patience=20):\n",
+    "    \"\"\"\n",
+    "    Full experiment runner with TensorBoard logging, early stopping, and diagnostics.\n",
+    "    model must return (logits, sphere_embeddings).\n",
+    "    \"\"\"\n",
+    "    writer = make_writer(experiment_id)\n",
+    "    criterion = nn.CrossEntropyLoss()\n",
+    "    optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=weight_decay)\n",
+    "    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs)\n",
+    "\n",
+    "    best_val_acc = 0.0\n",
+    "    best_epoch = 0\n",
+    "    no_improve = 0\n",
+    "\n",
+    "    # Log model info\n",
+    "    n_params = sum(p.numel() for p in model.parameters() if p.requires_grad)\n",
+    "    writer.add_text(\"config/params\", f\"Learned params: {n_params:,}\")\n",
+    "    writer.add_text(\"config/experiment\", experiment_id)\n",
+    "    print(f\"\\n{'='*70}\")\n",
+    "    print(f\"[EXP] {experiment_id}\")\n",
+    "    print(f\"[EXP] Learned params: {n_params:,}, Epochs: {epochs}, LR: {lr}\")\n",
+    "    print(f\"{'='*70}\")\n",
+    "\n",
+    "    for epoch in range(1, epochs + 1):\n",
+    "        t0 = time.time()\n",
+    "\n",
+    "        # Train\n",
+    "        train_metrics = train_one_epoch(\n",
+    "            model, train_loader, optimizer, criterion, device, epoch,\n",
+    "            extra_losses=extra_losses\n",
+    "        )\n",
+    "\n",
+    "        # Evaluate\n",
+    "        collect_emb = (epoch % log_embeddings_every == 0) or (epoch == epochs)\n",
+    "        val_result = evaluate(model, val_loader, criterion, device, collect_embeddings=collect_emb)\n",
+    "        val_metrics = val_result[\"metrics\"]\n",
+    "\n",
+    "        # Combine metrics\n",
+    "        all_metrics = {**train_metrics, **val_metrics}\n",
+    "        all_metrics[\"lr\"] = optimizer.param_groups[0][\"lr\"]\n",
+    "        all_metrics[\"epoch_time\"] = time.time() - t0\n",
+    "        all_metrics[\"train_val_gap\"] = train_metrics[\"accuracy/train\"] - val_metrics[\"accuracy/val\"]\n",
+    "\n",
+    "        # Constellation diagnostics\n",
+    "        if hasattr(model, \"constellation\"):\n",
+    "            constellation = model.constellation\n",
+    "        elif hasattr(model, \"classifier\") and hasattr(model.classifier, \"constellation\"):\n",
+    "            constellation = model.classifier.constellation\n",
+    "        else:\n",
+    "            constellation = None\n",
+    "\n",
+    "        if constellation is not None:\n",
+    "            log_anchor_histogram(writer, constellation, epoch)\n",
+    "\n",
+    "            # Push anchors periodically\n",
+    "            if epoch % push_anchors_every == 0 and \"embeddings\" in val_result:\n",
+    "                emb = val_result[\"embeddings\"][:2000].to(device)\n",
+    "                lab = torch.tensor(val_result[\"labels\"][:2000]).to(device)\n",
+    "                constellation.push_anchors(emb, lab)\n",
+    "                active, total_anch = constellation.anchor_activation(emb)\n",
+    "                all_metrics[\"anchors/active\"] = active\n",
+    "                all_metrics[\"anchors/total\"] = total_anch\n",
+    "                all_metrics[\"anchors/utilization\"] = active / total_anch\n",
+    "                print(f\"  [ANCHOR PUSH] Active: {active}/{total_anch} ({100*active/total_anch:.1f}%)\")\n",
+    "\n",
+    "        # Log all metrics\n",
+    "        log_epoch_metrics(writer, epoch, all_metrics)\n",
+    "        scheduler.step()\n",
+    "\n",
+    "        # Log embeddings visualization\n",
+    "        if collect_emb and \"embeddings\" in val_result:\n",
+    "            emb = val_result[\"embeddings\"][:2000]\n",
+    "            lab = val_result[\"labels\"][:2000]\n",
+    "            log_scatter_2d(writer, emb, lab, epoch, tag=\"val_embeddings\", method=\"pca\")\n",
+    "            if epoch == epochs:\n",
+    "                log_confusion_matrix(writer, val_result[\"preds\"], val_result[\"labels\"], epoch)\n",
+    "                log_scatter_2d(writer, emb, lab, epoch, tag=\"val_embeddings_tsne\", method=\"tsne\")\n",
+    "\n",
+    "        # Log triangulation stats from a sample batch\n",
+    "        if constellation is not None and collect_emb and \"embeddings\" in val_result:\n",
+    "            sample_emb = val_result[\"embeddings\"][:500].to(device)\n",
+    "            tri_sample = constellation(sample_emb)\n",
+    "            log_triangulation_stats(writer, tri_sample, epoch)\n",
+    "\n",
+    "        # Best model tracking\n",
+    "        val_acc = val_metrics[\"accuracy/val\"]\n",
+    "        if val_acc > best_val_acc:\n",
+    "            best_val_acc = val_acc\n",
+    "            best_epoch = epoch\n",
+    "            no_improve = 0\n",
+    "            writer.add_scalar(\"best/val_acc\", best_val_acc, epoch)\n",
+    "            writer.add_scalar(\"best/epoch\", best_epoch, epoch)\n",
+    "        else:\n",
+    "            no_improve += 1\n",
+    "\n",
+    "        # Print progress\n",
+    "        if epoch % 5 == 0 or epoch == 1 or epoch == epochs:\n",
+    "            print(f\"  E{epoch:3d} | TrainAcc {train_metrics['accuracy/train']:5.1f}% | \"\n",
+    "                  f\"ValAcc {val_acc:5.1f}% | Loss {all_metrics['loss/total']:.4f} | \"\n",
+    "                  f\"Best {best_val_acc:.1f}% @E{best_epoch} | \"\n",
+    "                  f\"Gap {all_metrics['train_val_gap']:+.1f}% | \"\n",
+    "                  f\"{all_metrics['epoch_time']:.1f}s\")\n",
+    "\n",
+    "        # Early stopping\n",
+    "        if no_improve >= patience:\n",
+    "            print(f\"  [EARLY STOP] No improvement for {patience} epochs. Best: {best_val_acc:.1f}% @E{best_epoch}\")\n",
+    "            break\n",
+    "\n",
+    "    writer.close()\n",
+    "    print(f\"\\n[RESULT] {experiment_id}: Best Val Acc = {best_val_acc:.2f}% at epoch {best_epoch}\")\n",
+    "    print(f\"[RESULT] Learned params: {n_params:,}\")\n",
+    "    return {\"experiment\": experiment_id, \"best_val_acc\": best_val_acc, \"best_epoch\": best_epoch,\n",
+    "            \"params\": n_params, \"final_train_acc\": train_metrics[\"accuracy/train\"]}\n",
+    "\n",
+    "# \u2500\u2500 Johnson-Lindenstrauss Projection \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
+    "class JLProjection(nn.Module):\n",
+    "    \"\"\"Fixed random Gaussian projection (not learned).\"\"\"\n",
+    "    def __init__(self, input_dim, output_dim):\n",
+    "        super().__init__()\n",
+    "        self.register_buffer(\"W\", torch.randn(input_dim, output_dim) / math.sqrt(output_dim))\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        return x @ self.W\n",
+    "\n",
+    "# \u2500\u2500 Results Scoreboard \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
+    "class Scoreboard:\n",
+    "    \"\"\"Track results across experiments in a session.\"\"\"\n",
+    "    def __init__(self):\n",
+    "        self.results = []\n",
+    "\n",
+    "    def add(self, result_dict):\n",
+    "        self.results.append(result_dict)\n",
+    "        self._print_table()\n",
+    "\n",
+    "    def _print_table(self):\n",
+    "        print(f\"\\n{'='*80}\")\n",
+    "        print(f\"{'SCOREBOARD':^80}\")\n",
+    "        print(f\"{'='*80}\")\n",
+    "        print(f\"{'Experiment':<40} {'Val Acc':>8} {'Params':>10} {'Best E':>7}\")\n",
+    "        print(f\"{'-'*40} {'-'*8} {'-'*10} {'-'*7}\")\n",
+    "        for r in sorted(self.results, key=lambda x: x[\"best_val_acc\"], reverse=True):\n",
+    "            print(f\"{r['experiment']:<40} {r['best_val_acc']:>7.2f}% {r['params']:>10,} {r['best_epoch']:>7}\")\n",
+    "        print(f\"{'='*80}\\n\")\n",
+    "\n",
+    "    def save(self, path=\"scoreboard.json\"):\n",
+    "        with open(path, \"w\") as f:\n",
+    "            json.dump(self.results, f, indent=2)\n",
+    "        print(f\"[SCOREBOARD] Saved to {path}\")\n",
+    "\n",
+    "scoreboard = Scoreboard()\n",
+    "print(\"[SHARED] All utilities loaded successfully.\")\n",
+    ""
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Additional Loss Functions"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "\n",
+    "from scipy.special import iv as bessel_i\n",
+    "\n",
+    "class vMFLoss(nn.Module):\n",
+    "    '''von Mises-Fisher based contrastive loss.'''\n",
+    "    def __init__(self, temperature=0.07, kappa_min=0.1, kappa_max=100.0):\n",
+    "        super().__init__()\n",
+    "        self.temperature = temperature\n",
+    "        self.kappa_min = kappa_min\n",
+    "        self.kappa_max = kappa_max\n",
+    "\n",
+    "    def forward(self, z, labels):\n",
+    "        z = l2_normalize(z)\n",
+    "        B = z.shape[0]\n",
+    "\n",
+    "        # Compute similarity\n",
+    "        sim = torch.mm(z, z.T)  # (B, B)\n",
+    "\n",
+    "        # Estimate kappa from concentration of similarities\n",
+    "        sim_mean = sim.mean().item()\n",
+    "        kappa = np.clip(sim_mean * 50, self.kappa_min, self.kappa_max)\n",
+    "\n",
+    "        # vMF log-likelihood (approximation)\n",
+    "        # log p(z_i | mu) ~ kappa * z_i^T mu - log(C_d(kappa))\n",
+    "        label_mask = (labels.unsqueeze(0) == labels.unsqueeze(1)).float()\n",
+    "        label_mask.fill_diagonal_(0)\n",
+    "\n",
+    "        # Positive pairs have higher kappa\n",
+    "        log_likelihood = kappa * sim\n",
+    "        pos_mask = label_mask\n",
+    "        neg_mask = 1 - label_mask\n",
+    "        neg_mask.fill_diagonal_(0)\n",
+    "\n",
+    "        pos_ll = (log_likelihood * pos_mask).sum(dim=1)\n",
+    "        neg_ll = (log_likelihood * neg_mask).sum(dim=1)\n",
+    "\n",
+    "        # Contrastive objective\n",
+    "        loss = -torch.log(torch.exp(pos_ll) / (torch.exp(pos_ll) + torch.exp(neg_ll) + 1e-8) + 1e-8).mean()\n",
+    "        return loss\n",
+    "\n",
+    "class FisherRaoLoss(nn.Module):\n",
+    "    '''Fisher-Rao distance as loss: d_FR(p,q) = 2\u00b7arccos(\u222b\u221a(pq)).'''\n",
+    "    def __init__(self, temperature=0.07):\n",
+    "        super().__init__()\n",
+    "        self.temperature = temperature\n",
+    "\n",
+    "    def forward(self, z, labels):\n",
+    "        z = l2_normalize(z)\n",
+    "        B = z.shape[0]\n",
+    "\n",
+    "        # Pairwise cosine similarity\n",
+    "        sim = torch.mm(z, z.T)  # (B, B)\n",
+    "\n",
+    "        # Fisher-Rao distance ~ arccos(similarity)\n",
+    "        # For normalized vectors, Bhattacharyya coefficient \u2248 sim\n",
+    "        dist = torch.acos(torch.clamp(sim, -1 + 1e-7, 1 - 1e-7))\n",
+    "\n",
+    "        label_mask = (labels.unsqueeze(0) == labels.unsqueeze(1)).float()\n",
+    "        label_mask.fill_diagonal_(0)\n",
+    "\n",
+    "        # Minimize distance within class, maximize across class\n",
+    "        pos_dist = (dist * label_mask).sum(dim=1) / (label_mask.sum(dim=1) + 1e-8)\n",
+    "        neg_dist = (dist * (1 - label_mask)).sum(dim=1) / ((1 - label_mask).sum(dim=1) + 1e-8)\n",
+    "\n",
+    "        loss = (pos_dist - neg_dist + 1.0).relu().mean()\n",
+    "        return loss\n",
+    "\n",
+    "class SlicedWassersteinLoss(nn.Module):\n",
+    "    '''Sliced Wasserstein distance on hypersphere.'''\n",
+    "    def __init__(self, n_projections=50):\n",
+    "        super().__init__()\n",
+    "        self.n_projections = n_projections\n",
+    "\n",
+    "    def forward(self, z, labels):\n",
+    "        z = l2_normalize(z)\n",
+    "        B, d = z.shape\n",
+    "\n",
+    "        # Random projections on sphere\n",
+    "        thetas = torch.randn(self.n_projections, d, device=z.device)\n",
+    "        thetas = F.normalize(thetas, dim=1)\n",
+    "\n",
+    "        # Project embeddings\n",
+    "        proj_z = torch.mm(z, thetas.T)  # (B, n_proj)\n",
+    "\n",
+    "        label_mask = (labels.unsqueeze(0) == labels.unsqueeze(1)).float()\n",
+    "        label_mask.fill_diagonal_(0)\n",
+    "\n",
+    "        # Sort projections\n",
+    "        proj_sorted, _ = torch.sort(proj_z, dim=0)\n",
+    "\n",
+    "        # Compute sliced Wasserstein\n",
+    "        loss = 0.0\n",
+    "        for b in range(B):\n",
+    "            pos_idx = (labels == labels[b]).nonzero(as_tuple=True)[0]\n",
+    "            if len(pos_idx) > 1:\n",
+    "                pos_proj = proj_z[pos_idx]\n",
+    "                neg_proj = proj_z[labels != labels[b]]\n",
+    "\n",
+    "                if len(neg_proj) > 0:\n",
+    "                    for i in range(self.n_projections):\n",
+    "                        w_dist = torch.abs(pos_proj[:, i].sort()[0] - neg_proj[:, i].sort()[0]).mean()\n",
+    "                        loss = loss + w_dist\n",
+    "\n",
+    "        return loss / (B + 1e-8)\n",
+    "\n",
+    "class GeometricAutogradLoss(nn.Module):\n",
+    "    '''Geometric loss: tangential projection + separation preservation.'''\n",
+    "    def __init__(self):\n",
+    "        super().__init__()\n",
+    "\n",
+    "    def forward(self, z, labels):\n",
+    "        z = l2_normalize(z)\n",
+    "        B = z.shape[0]\n",
+    "\n",
+    "        # Compute gradients manually if needed\n",
+    "        z.requires_grad_(True)\n",
+    "\n",
+    "        # Project to tangent space at each point\n",
+    "        # T_z S^d = {v : v^T z = 0}\n",
+    "        I = torch.eye(z.shape[1], device=z.device)\n",
+    "        z_outer = torch.bmm(z.unsqueeze(2), z.unsqueeze(1))  # (B, d, d)\n",
+    "        proj = I.unsqueeze(0) - z_outer  # (B, d, d)\n",
+    "\n",
+    "        # Separation: maximize distance between different classes\n",
+    "        label_mask = (labels.unsqueeze(0) == labels.unsqueeze(1)).float()\n",
+    "        label_mask.fill_diagonal_(0)\n",
+    "\n",
+    "        # Pairwise distances\n",
+    "        z_diff = z.unsqueeze(1) - z.unsqueeze(0)  # (B, B, d)\n",
+    "        distances = torch.norm(z_diff, dim=2)  # (B, B)\n",
+    "\n",
+    "        # Loss: push different classes apart\n",
+    "        loss = (distances * (1 - label_mask)).mean()\n",
+    "        return loss\n",
+    "\n",
+    "print(\"[LOSSES] All additional losses loaded.\")\n",
+    ""
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 6.2a: CE-Only Baseline (No Contrastive Loss)\n",
+    "\n",
+    "Measure ceiling with scattering front-end"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "\n",
+    "from kymatio.torch import Scattering2D\n",
+    "\n",
+    "print(\"[6.2a] CE-Only Baseline\")\n",
+    "print(\"=\"*70)\n",
+    "\n",
+    "class CEOnlyModel(nn.Module):\n",
+    "    '''Scattering2D \u2192 JL(128) \u2192 L2 norm \u2192 constellation \u2192 classify (CE only).'''\n",
+    "    def __init__(self, embed_dim=128, num_classes=10):\n",
+    "        super().__init__()\n",
+    "        self.scattering = Scattering2D(J=2, L=8, shape=(32, 32), max_order=2)\n",
+    "        self.scatt_dim = 81  # J=2, L=8 default output\n",
+    "        self.jl = JLProjection(self.scatt_dim, 128)\n",
+    "        self.classifier = ConstellationClassifier(embed_dim, num_anchors=64, num_classes=num_classes)\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        scatt = self.scattering(x)  # (B, 81, 32, 32)\n",
+    "        scatt_flat = scatt.view(scatt.shape[0], -1)  # (B, 81*32*32)\n",
+    "\n",
+    "        # Dimensionality reduction\n",
+    "        reduced = F.adaptive_avg_pool2d(scatt.view(-1, 81, 32, 32), (1, 1)).squeeze(-1).squeeze(-1)  # (B, 81)\n",
+    "        jl_feat = self.jl(reduced)  # (B, 128)\n",
+    "        emb = l2_normalize(jl_feat)\n",
+    "\n",
+    "        logits, tri = self.classifier(emb)\n",
+    "        return logits, emb\n",
+    "\n",
+    "model = CEOnlyModel(embed_dim=128, num_classes=10).to(device)\n",
+    "print(f\"[6.2a] Model params: {sum(p.numel() for p in model.parameters()):,}\")\n",
+    "\n",
+    "train_loader, val_loader = get_cifar10_loaders(batch_size=256, augment=True)\n",
+    "\n",
+    "result_6_2a = run_experiment(\n",
+    "    \"6.2a_CE_Only\",\n",
+    "    model, train_loader, val_loader, device,\n",
+    "    epochs=100, lr=1e-3, weight_decay=1e-4, patience=20\n",
+    ")\n",
+    "scoreboard.add(result_6_2a)\n",
+    ""
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 6.2b: Class-Conditional InfoNCE\n",
+    "\n",
+    "Positive = same class, not same image"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "\n",
+    "print(\"[6.2b] Class-Conditional InfoNCE\")\n",
+    "print(\"=\"*70)\n",
+    "\n",
+    "class ClassConditionalModel(nn.Module):\n",
+    "    '''CE + Class-Conditional InfoNCE loss.'''\n",
+    "    def __init__(self, embed_dim=128, num_classes=10):\n",
+    "        super().__init__()\n",
+    "        self.scattering = Scattering2D(J=2, L=8, shape=(32, 32), max_order=2)\n",
+    "        self.jl = JLProjection(81, 128)\n",
+    "        self.classifier = ConstellationClassifier(embed_dim, num_anchors=64, num_classes=num_classes)\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        scatt = self.scattering(x)\n",
+    "        reduced = F.adaptive_avg_pool2d(scatt.view(-1, 81, 32, 32), (1, 1)).squeeze(-1).squeeze(-1)\n",
+    "        jl_feat = self.jl(reduced)\n",
+    "        emb = l2_normalize(jl_feat)\n",
+    "        logits, tri = self.classifier(emb)\n",
+    "        return logits, emb\n",
+    "\n",
+    "model = ClassConditionalModel(embed_dim=128, num_classes=10).to(device)\n",
+    "print(f\"[6.2b] Model params: {sum(p.numel() for p in model.parameters()):,}\")\n",
+    "\n",
+    "class_info_loss = ClassConditionalInfoNCE(temperature=0.07)\n",
+    "\n",
+    "result_6_2b = run_experiment(\n",
+    "    \"6.2b_ClassConditionalInfoNCE\",\n",
+    "    model, train_loader, val_loader, device,\n",
+    "    epochs=100, lr=1e-3, weight_decay=1e-4,\n",
+    "    extra_losses=[(\"InfoNCE_class\", class_info_loss, 0.5)],\n",
+    "    patience=20\n",
+    ")\n",
+    "scoreboard.add(result_6_2b)\n",
+    ""
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 6.2c: vMF-based Contrastive Loss\n",
+    "\n",
+    "Replace dot-product similarity with vMF log-likelihood"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "\n",
+    "print(\"[6.2c] vMF-based Contrastive Loss\")\n",
+    "print(\"=\"*70)\n",
+    "\n",
+    "model = ClassConditionalModel(embed_dim=128, num_classes=10).to(device)\n",
+    "print(f\"[6.2c] Model params: {sum(p.numel() for p in model.parameters()):,}\")\n",
+    "\n",
+    "vmf_loss = vMFLoss(temperature=0.07)\n",
+    "\n",
+    "result_6_2c = run_experiment(\n",
+    "    \"6.2c_vMF_Contrastive\",\n",
+    "    model, train_loader, val_loader, device,\n",
+    "    epochs=100, lr=1e-3, weight_decay=1e-4,\n",
+    "    extra_losses=[(\"vMF\", vmf_loss, 0.5)],\n",
+    "    patience=20\n",
+    ")\n",
+    "scoreboard.add(result_6_2c)\n",
+    ""
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 6.2d: Fisher-Rao Distance as Loss\n",
+    "\n",
+    "Geometric distance d_FR(p,q) = 2\u00b7arccos(\u222b\u221a(pq))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "\n",
+    "print(\"[6.2d] Fisher-Rao Distance Loss\")\n",
+    "print(\"=\"*70)\n",
+    "\n",
+    "model = ClassConditionalModel(embed_dim=128, num_classes=10).to(device)\n",
+    "print(f\"[6.2d] Model params: {sum(p.numel() for p in model.parameters()):,}\")\n",
+    "\n",
+    "fr_loss = FisherRaoLoss(temperature=0.07)\n",
+    "\n",
+    "result_6_2d = run_experiment(\n",
+    "    \"6.2d_FisherRao\",\n",
+    "    model, train_loader, val_loader, device,\n",
+    "    epochs=100, lr=1e-3, weight_decay=1e-4,\n",
+    "    extra_losses=[(\"FisherRao\", fr_loss, 0.5)],\n",
+    "    patience=20\n",
+    ")\n",
+    "scoreboard.add(result_6_2d)\n",
+    ""
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 6.2e: Sliced Spherical Wasserstein Distance\n",
+    "\n",
+    "Distribution matching loss on hypersphere"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "\n",
+    "print(\"[6.2e] Sliced Spherical Wasserstein Distance\")\n",
+    "print(\"=\"*70)\n",
+    "\n",
+    "model = ClassConditionalModel(embed_dim=128, num_classes=10).to(device)\n",
+    "print(f\"[6.2e] Model params: {sum(p.numel() for p in model.parameters()):,}\")\n",
+    "\n",
+    "swd_loss = SlicedWassersteinLoss(n_projections=50)\n",
+    "\n",
+    "result_6_2e = run_experiment(\n",
+    "    \"6.2e_SlicedWasserstein\",\n",
+    "    model, train_loader, val_loader, device,\n",
+    "    epochs=100, lr=1e-3, weight_decay=1e-4,\n",
+    "    extra_losses=[(\"SWD\", swd_loss, 0.3)],\n",
+    "    patience=20\n",
+    ")\n",
+    "scoreboard.add(result_6_2e)\n",
+    ""
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 6.2f: Geometric Autograd Loss\n",
+    "\n",
+    "Tangential projection + separation preservation on gradients"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "\n",
+    "print(\"[6.2f] Geometric Autograd Loss\")\n",
+    "print(\"=\"*70)\n",
+    "\n",
+    "model = ClassConditionalModel(embed_dim=128, num_classes=10).to(device)\n",
+    "print(f\"[6.2f] Model params: {sum(p.numel() for p in model.parameters()):,}\")\n",
+    "\n",
+    "geom_loss = GeometricAutogradLoss()\n",
+    "\n",
+    "result_6_2f = run_experiment(\n",
+    "    \"6.2f_GeometricAutograd\",\n",
+    "    model, train_loader, val_loader, device,\n",
+    "    epochs=100, lr=1e-3, weight_decay=1e-4,\n",
+    "    extra_losses=[(\"Geometric\", geom_loss, 0.3)],\n",
+    "    patience=20\n",
+    ")\n",
+    "scoreboard.add(result_6_2f)\n",
+    ""
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 6.3a: Anchor Push Frequency Sweep\n",
+    "\n",
+    "Every 10, 25, 50, 100, 200 batches"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "\n",
+    "print(\"[6.3a] Anchor Push Frequency Sweep\")\n",
+    "print(\"=\"*70)\n",
+    "\n",
+    "for freq in [10, 25, 50, 100, 200]:\n",
+    "    model = ClassConditionalModel(embed_dim=128, num_classes=10).to(device)\n",
+    "\n",
+    "    # Calculate push_anchors_every from batch count\n",
+    "    # ~196 batches per epoch (50k / 256)\n",
+    "    push_every_epochs = max(1, freq // 196)\n",
+    "\n",
+    "    result = run_experiment(\n",
+    "        f\"6.3a_AnchorPushFreq_{freq}\",\n",
+    "        model, train_loader, val_loader, device,\n",
+    "        epochs=100, lr=1e-3, weight_decay=1e-4,\n",
+    "        push_anchors_every=push_every_epochs,\n",
+    "        patience=20\n",
+    "    )\n",
+    "    scoreboard.add(result)\n",
+    "    print(f\"  Frequency: {freq} batches \u2192 push_anchors_every={push_every_epochs} epochs\")\n",
+    ""
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 6.3b: Anchor Push with vMF-weighted Centroids"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "\n",
+    "print(\"[6.3b] Anchor Push with vMF-weighted Centroids\")\n",
+    "print(\"=\"*70)\n",
+    "\n",
+    "class ConstellationWithvMFPush(nn.Module):\n",
+    "    '''Constellation with vMF-weighted anchor push.'''\n",
+    "    def __init__(self, embed_dim, num_anchors=64):\n",
+    "        super().__init__()\n",
+    "        self.embed_dim = embed_dim\n",
+    "        self.num_anchors = num_anchors\n",
+    "        anchors = self._init_anchors(embed_dim, num_anchors, \"uniform\")\n",
+    "        self.anchors = nn.Parameter(anchors, requires_grad=False)\n",
+    "        self.kappa = 10.0  # von Mises-Fisher concentration\n",
+    "\n",
+    "    def _init_anchors(self, d, K, method):\n",
+    "        A = torch.randn(K, d)\n",
+    "        Q, _ = torch.linalg.qr(A.T)\n",
+    "        anchors = Q.T[:K]\n",
+    "        if K > d:\n",
+    "            extra = F.normalize(torch.randn(K - d, d), dim=-1)\n",
+    "            anchors = torch.cat([anchors, extra], dim=0)\n",
+    "        return F.normalize(anchors, dim=-1)\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        x_norm = l2_normalize(x)\n",
+    "        cos_sim = torch.mm(x_norm, self.anchors.T)\n",
+    "        distances = 1.0 - cos_sim\n",
+    "        return distances\n",
+    "\n",
+    "    def push_anchors_vmf(self, embeddings, labels, num_classes=10, momentum=0.1):\n",
+    "        '''Push anchors using vMF-weighted centroids.'''\n",
+    "        with torch.no_grad():\n",
+    "            for c in range(num_classes):\n",
+    "                mask = (labels == c)\n",
+    "                if mask.sum() < 2:\n",
+    "                    continue\n",
+    "\n",
+    "                emb_class = embeddings[mask]\n",
+    "                emb_norm = l2_normalize(emb_class)\n",
+    "\n",
+    "                # vMF weights: higher for points closer to class mean\n",
+    "                mean_dir = l2_normalize(emb_norm.mean(dim=0, keepdim=True))\n",
+    "                sim_to_mean = torch.mm(emb_norm, mean_dir.T).squeeze()\n",
+    "                weights = torch.exp(self.kappa * sim_to_mean)\n",
+    "                weights = weights / (weights.sum() + 1e-8)\n",
+    "\n",
+    "                # Weighted centroid\n",
+    "                centroid = (emb_norm * weights.unsqueeze(1)).sum(dim=0, keepdim=True)\n",
+    "                centroid = l2_normalize(centroid)\n",
+    "\n",
+    "                start = c * (self.num_anchors // num_classes)\n",
+    "                end = start + (self.num_anchors // num_classes)\n",
+    "                self.anchors.data[start:end] = l2_normalize(\n",
+    "                    (1 - momentum) * self.anchors.data[start:end] + momentum * centroid\n",
+    "                )\n",
+    "\n",
+    "class vMFAnchorPushModel(nn.Module):\n",
+    "    def __init__(self, embed_dim=128, num_classes=10):\n",
+    "        super().__init__()\n",
+    "        self.scattering = Scattering2D(J=2, L=8, shape=(32, 32), max_order=2)\n",
+    "        self.jl = JLProjection(81, 128)\n",
+    "        self.constellation = ConstellationWithvMFPush(embed_dim, num_anchors=64)\n",
+    "        self.patchwork = Patchwork(64, 2)\n",
+    "        self.head = nn.Linear(64, num_classes)\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        scatt = self.scattering(x)\n",
+    "        reduced = F.adaptive_avg_pool2d(scatt.view(-1, 81, 32, 32), (1, 1)).squeeze(-1).squeeze(-1)\n",
+    "        jl_feat = self.jl(reduced)\n",
+    "        emb = l2_normalize(jl_feat)\n",
+    "        tri = self.constellation(emb)\n",
+    "        pw = self.patchwork(tri)\n",
+    "        logits = self.head(pw)\n",
+    "        return logits, emb\n",
+    "\n",
+    "model = vMFAnchorPushModel(embed_dim=128, num_classes=10).to(device)\n",
+    "print(f\"[6.3b] Model params: {sum(p.numel() for p in model.parameters()):,}\")\n",
+    "\n",
+    "result_6_3b = run_experiment(\n",
+    "    \"6.3b_vMFAnchorPush\",\n",
+    "    model, train_loader, val_loader, device,\n",
+    "    epochs=100, lr=1e-3, weight_decay=1e-4,\n",
+    "    push_anchors_every=50,\n",
+    "    patience=20\n",
+    ")\n",
+    "scoreboard.add(result_6_3b)\n",
+    ""
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 6.3c: Anchor Birth/Death\n",
+    "\n",
+    "Add anchors where density high, remove where unused"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "\n",
+    "print(\"[6.3c] Anchor Birth/Death\")\n",
+    "print(\"=\"*70)\n",
+    "\n",
+    "class DynamicConstellation(nn.Module):\n",
+    "    '''Constellation with birth/death of anchors based on density.'''\n",
+    "    def __init__(self, embed_dim, num_anchors=64):\n",
+    "        super().__init__()\n",
+    "        self.embed_dim = embed_dim\n",
+    "        self.num_anchors = num_anchors\n",
+    "        anchors = self._init_anchors(embed_dim, num_anchors, \"uniform\")\n",
+    "        self.anchors = nn.Parameter(anchors, requires_grad=False)\n",
+    "        self.anchor_activity = torch.zeros(num_anchors)\n",
+    "\n",
+    "    def _init_anchors(self, d, K, method):\n",
+    "        A = torch.randn(K, d)\n",
+    "        Q, _ = torch.linalg.qr(A.T)\n",
+    "        anchors = Q.T[:K]\n",
+    "        if K > d:\n",
+    "            extra = F.normalize(torch.randn(K - d, d), dim=-1)\n",
+    "            anchors = torch.cat([anchors, extra], dim=0)\n",
+    "        return F.normalize(anchors, dim=-1)\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        x_norm = l2_normalize(x)\n",
+    "        cos_sim = torch.mm(x_norm, self.anchors.T)\n",
+    "        distances = 1.0 - cos_sim\n",
+    "\n",
+    "        # Track activity\n",
+    "        nearest = cos_sim.argmax(dim=1)\n",
+    "        for anc_idx in range(self.num_anchors):\n",
+    "            self.anchor_activity[anc_idx] = (nearest == anc_idx).sum().float()\n",
+    "\n",
+    "        return distances\n",
+    "\n",
+    "    def birth_death(self, embeddings, labels, num_classes=10, density_threshold=0.1):\n",
+    "        '''Add anchors in high-density regions, remove unused ones.'''\n",
+    "        with torch.no_grad():\n",
+    "            emb_norm = l2_normalize(embeddings)\n",
+    "\n",
+    "            # Find high-density regions\n",
+    "            cos_sim = torch.mm(emb_norm, emb_norm.T)\n",
+    "            density = cos_sim.mean(dim=1)\n",
+    "            high_density_idx = torch.where(density > density_threshold)[0]\n",
+    "\n",
+    "            # Find unused anchors\n",
+    "            unused_idx = torch.where(self.anchor_activity < 1)[0]\n",
+    "\n",
+    "            # Birth: add new anchors at high-density centroids\n",
+    "            if len(high_density_idx) > 0 and len(unused_idx) > 0:\n",
+    "                new_anchor = l2_normalize(emb_norm[high_density_idx].mean(dim=0, keepdim=True))\n",
+    "                self.anchors.data[unused_idx[0]] = new_anchor[0]\n",
+    "\n",
+    "            # Death: replace unused anchors\n",
+    "            if len(unused_idx) > 0:\n",
+    "                for anc_idx in unused_idx[:len(unused_idx)//2]:\n",
+    "                    random_idx = torch.randint(0, len(embeddings), (1,)).item()\n",
+    "                    self.anchors.data[anc_idx] = l2_normalize(embeddings[random_idx:random_idx+1])\n",
+    "\n",
+    "class BirthDeathModel(nn.Module):\n",
+    "    def __init__(self, embed_dim=128, num_classes=10):\n",
+    "        super().__init__()\n",
+    "        self.scattering = Scattering2D(J=2, L=8, shape=(32, 32), max_order=2)\n",
+    "        self.jl = JLProjection(81, 128)\n",
+    "        self.constellation = DynamicConstellation(embed_dim, num_anchors=64)\n",
+    "        self.patchwork = Patchwork(64, 2)\n",
+    "        self.head = nn.Linear(64, num_classes)\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        scatt = self.scattering(x)\n",
+    "        reduced = F.adaptive_avg_pool2d(scatt.view(-1, 81, 32, 32), (1, 1)).squeeze(-1).squeeze(-1)\n",
+    "        jl_feat = self.jl(reduced)\n",
+    "        emb = l2_normalize(jl_feat)\n",
+    "        tri = self.constellation(emb)\n",
+    "        pw = self.patchwork(tri)\n",
+    "        logits = self.head(pw)\n",
+    "        return logits, emb\n",
+    "\n",
+    "model = BirthDeathModel(embed_dim=128, num_classes=10).to(device)\n",
+    "print(f\"[6.3c] Model params: {sum(p.numel() for p in model.parameters()):,}\")\n",
+    "\n",
+    "result_6_3c = run_experiment(\n",
+    "    \"6.3c_AnchorBirthDeath\",\n",
+    "    model, train_loader, val_loader, device,\n",
+    "    epochs=100, lr=1e-3, weight_decay=1e-4,\n",
+    "    push_anchors_every=25,\n",
+    "    patience=20\n",
+    ")\n",
+    "scoreboard.add(result_6_3c)\n",
+    ""
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 6.3d: Anchor Dropout Sweep\n",
+    "\n",
+    "0%, 5%, 15%, 30%, 50%"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "\n",
+    "print(\"[6.3d] Anchor Dropout Sweep\")\n",
+    "print(\"=\"*70)\n",
+    "\n",
+    "class DropoutConstellation(nn.Module):\n",
+    "    def __init__(self, embed_dim, num_anchors=64, dropout_rate=0.0):\n",
+    "        super().__init__()\n",
+    "        self.embed_dim = embed_dim\n",
+    "        self.num_anchors = num_anchors\n",
+    "        self.dropout_rate = dropout_rate\n",
+    "        anchors = self._init_anchors(embed_dim, num_anchors, \"uniform\")\n",
+    "        self.anchors = nn.Parameter(anchors, requires_grad=False)\n",
+    "\n",
+    "    def _init_anchors(self, d, K, method):\n",
+    "        A = torch.randn(K, d)\n",
+    "        Q, _ = torch.linalg.qr(A.T)\n",
+    "        anchors = Q.T[:K]\n",
+    "        if K > d:\n",
+    "            extra = F.normalize(torch.randn(K - d, d), dim=-1)\n",
+    "            anchors = torch.cat([anchors, extra], dim=0)\n",
+    "        return F.normalize(anchors, dim=-1)\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        x_norm = l2_normalize(x)\n",
+    "\n",
+    "        # Apply dropout to anchors\n",
+    "        if self.training and self.dropout_rate > 0:\n",
+    "            mask = torch.bernoulli(torch.ones(self.num_anchors, device=x.device) * (1 - self.dropout_rate))\n",
+    "            anch_active = self.anchors * mask.unsqueeze(1)\n",
+    "            anch_active = F.normalize(anch_active, dim=1)\n",
+    "        else:\n",
+    "            anch_active = self.anchors\n",
+    "\n",
+    "        cos_sim = torch.mm(x_norm, anch_active.T)\n",
+    "        distances = 1.0 - cos_sim\n",
+    "        return distances\n",
+    "\n",
+    "class DropoutModel(nn.Module):\n",
+    "    def __init__(self, embed_dim=128, num_classes=10, dropout_rate=0.0):\n",
+    "        super().__init__()\n",
+    "        self.scattering = Scattering2D(J=2, L=8, shape=(32, 32), max_order=2)\n",
+    "        self.jl = JLProjection(81, 128)\n",
+    "        self.constellation = DropoutConstellation(embed_dim, num_anchors=64, dropout_rate=dropout_rate)\n",
+    "        self.patchwork = Patchwork(64, 2)\n",
+    "        self.head = nn.Linear(64, num_classes)\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        scatt = self.scattering(x)\n",
+    "        reduced = F.adaptive_avg_pool2d(scatt.view(-1, 81, 32, 32), (1, 1)).squeeze(-1).squeeze(-1)\n",
+    "        jl_feat = self.jl(reduced)\n",
+    "        emb = l2_normalize(jl_feat)\n",
+    "        tri = self.constellation(emb)\n",
+    "        pw = self.patchwork(tri)\n",
+    "        logits = self.head(pw)\n",
+    "        return logits, emb\n",
+    "\n",
+    "for dropout_rate in [0.0, 0.05, 0.15, 0.30, 0.50]:\n",
+    "    model = DropoutModel(embed_dim=128, num_classes=10, dropout_rate=dropout_rate).to(device)\n",
+    "    result = run_experiment(\n",
+    "        f\"6.3d_AnchorDropout_{int(dropout_rate*100)}pct\",\n",
+    "        model, train_loader, val_loader, device,\n",
+    "        epochs=100, lr=1e-3, weight_decay=1e-4,\n",
+    "        push_anchors_every=50,\n",
+    "        patience=20\n",
+    "    )\n",
+    "    scoreboard.add(result)\n",
+    "    print(f\"  Dropout rate: {dropout_rate*100:.0f}%\")\n",
+    ""
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Final Scoreboard"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "\n",
+    "print(\"\\n\" + \"=\"*80)\n",
+    "print(\"EXPERIMENT 6 COMPLETE: Information-Theoretic Losses & Anchor Management\")\n",
+    "print(\"=\"*80)\n",
+    "scoreboard._print_table()\n",
+    "scoreboard.save(\"scoreboard_exp6.json\")\n",
+    ""
+   ],
+   "execution_count": null,
+   "outputs": []
+  }
+ ]
+}

spectral/experiment_7_composite_pipelines.ipynb

@@ -0,0 +1,1381 @@
+{
+ "nbformat": 4,
+ "nbformat_minor": 5,
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.10.0"
+  },
+  "accelerator": "GPU",
+  "colab": {
+   "provenance": [],
+   "gpuType": "T4"
+  }
+ },
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Experiment 7: Composite Pipeline Tests\n",
+    "\n",
+    "Reference pipeline, multi-signature fusion, hierarchical constellation."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "\n",
+    "!pip install -q kymatio scipy tensorboard scikit-learn matplotlib numpy torch torchvision torchaudio\n",
+    ""
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Shared Utilities"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "\n",
+    "# \u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\n",
+    "# SHARED UTILITIES \u2014 GeoLIP Hypersphere Experiments\n",
+    "# \u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\n",
+    "import os, sys, time, math, json, datetime, warnings\n",
+    "import numpy as np\n",
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "import torch.nn.functional as F\n",
+    "from torch.utils.data import DataLoader\n",
+    "from torch.utils.tensorboard import SummaryWriter\n",
+    "import torchvision\n",
+    "import torchvision.transforms as T\n",
+    "import matplotlib\n",
+    "matplotlib.use(\"Agg\")\n",
+    "import matplotlib.pyplot as plt\n",
+    "from collections import defaultdict\n",
+    "\n",
+    "warnings.filterwarnings(\"ignore\", category=UserWarning)\n",
+    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
+    "print(f\"[DEVICE] {device}\" + (f\" \u2014 {torch.cuda.get_device_name()}\" if device.type == \"cuda\" else \"\"))\n",
+    "\n",
+    "# \u2500\u2500 CIFAR-10 Loaders \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
+    "CIFAR10_MEAN = (0.4914, 0.4822, 0.4465)\n",
+    "CIFAR10_STD  = (0.2470, 0.2435, 0.2616)\n",
+    "CIFAR10_CLASSES = (\"plane\",\"car\",\"bird\",\"cat\",\"deer\",\"dog\",\"frog\",\"horse\",\"ship\",\"truck\")\n",
+    "\n",
+    "def get_cifar10_loaders(batch_size=256, num_workers=2, augment=True):\n",
+    "    \"\"\"Return train_loader, val_loader for CIFAR-10.\"\"\"\n",
+    "    tf_train = T.Compose([\n",
+    "        T.RandomCrop(32, padding=4),\n",
+    "        T.RandomHorizontalFlip(),\n",
+    "        T.ToTensor(),\n",
+    "        T.Normalize(CIFAR10_MEAN, CIFAR10_STD),\n",
+    "    ]) if augment else T.Compose([T.ToTensor(), T.Normalize(CIFAR10_MEAN, CIFAR10_STD)])\n",
+    "    tf_val = T.Compose([T.ToTensor(), T.Normalize(CIFAR10_MEAN, CIFAR10_STD)])\n",
+    "\n",
+    "    train_ds = torchvision.datasets.CIFAR10(root=\"./data\", train=True,  download=True, transform=tf_train)\n",
+    "    val_ds   = torchvision.datasets.CIFAR10(root=\"./data\", train=False, download=True, transform=tf_val)\n",
+    "    train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True,  num_workers=num_workers, pin_memory=True, drop_last=True)\n",
+    "    val_loader   = DataLoader(val_ds,   batch_size=batch_size, shuffle=False, num_workers=num_workers, pin_memory=True)\n",
+    "    print(f\"[DATA] CIFAR-10: {len(train_ds)} train, {len(val_ds)} val, batch_size={batch_size}\")\n",
+    "    return train_loader, val_loader\n",
+    "\n",
+    "def get_cifar10_raw_loaders(batch_size=256, num_workers=2):\n",
+    "    \"\"\"Return loaders with only ToTensor (no normalization) for custom front-ends.\"\"\"\n",
+    "    tf = T.Compose([T.ToTensor()])\n",
+    "    train_ds = torchvision.datasets.CIFAR10(root=\"./data\", train=True,  download=True, transform=tf)\n",
+    "    val_ds   = torchvision.datasets.CIFAR10(root=\"./data\", train=False, download=True, transform=tf)\n",
+    "    train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True,  num_workers=num_workers, pin_memory=True, drop_last=True)\n",
+    "    val_loader   = DataLoader(val_ds,   batch_size=batch_size, shuffle=False, num_workers=num_workers, pin_memory=True)\n",
+    "    print(f\"[DATA] CIFAR-10 (raw): {len(train_ds)} train, {len(val_ds)} val\")\n",
+    "    return train_loader, val_loader\n",
+    "\n",
+    "# \u2500\u2500 L2 Normalization to Hypersphere \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
+    "def l2_normalize(x, dim=-1, eps=1e-8):\n",
+    "    \"\"\"Project to unit hypersphere S^{d-1}.\"\"\"\n",
+    "    return F.normalize(x, p=2, dim=dim, eps=eps)\n",
+    "\n",
+    "# \u2500\u2500 Constellation Module \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
+    "class Constellation(nn.Module):\n",
+    "    \"\"\"\n",
+    "    K anchors on S^{d-1}. Triangulation: compute distance from input to each anchor.\n",
+    "    Output: K-dim vector of (1 - cos_sim) distances = triangulation profile.\n",
+    "    \"\"\"\n",
+    "    def __init__(self, embed_dim, num_anchors=64, init=\"uniform\"):\n",
+    "        super().__init__()\n",
+    "        self.embed_dim = embed_dim\n",
+    "        self.num_anchors = num_anchors\n",
+    "        anchors = self._init_anchors(embed_dim, num_anchors, init)\n",
+    "        self.anchors = nn.Parameter(anchors, requires_grad=False)  # anchors are NOT learned\n",
+    "\n",
+    "    def _init_anchors(self, d, K, method):\n",
+    "        if method == \"uniform\":\n",
+    "            # QR-based uniform hypersphere initialization\n",
+    "            A = torch.randn(K, d)\n",
+    "            Q, _ = torch.linalg.qr(A.T)\n",
+    "            anchors = Q.T[:K]\n",
+    "            # If K > d, fill remaining with random unit vectors\n",
+    "            if K > d:\n",
+    "                extra = F.normalize(torch.randn(K - d, d), dim=-1)\n",
+    "                anchors = torch.cat([anchors, extra], dim=0)\n",
+    "            return F.normalize(anchors, dim=-1)\n",
+    "        elif method == \"random\":\n",
+    "            return F.normalize(torch.randn(K, d), dim=-1)\n",
+    "        else:\n",
+    "            return F.normalize(torch.randn(K, d), dim=-1)\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        \"\"\"x: (B, d) on S^{d-1} \u2192 (B, K) triangulation distances.\"\"\"\n",
+    "        x_norm = l2_normalize(x)\n",
+    "        cos_sim = torch.mm(x_norm, self.anchors.T)  # (B, K)\n",
+    "        distances = 1.0 - cos_sim  # in [0, 2]\n",
+    "        return distances\n",
+    "\n",
+    "    def push_anchors(self, embeddings, labels, num_classes=10, momentum=0.1):\n",
+    "        \"\"\"Push anchors toward class centroids (call periodically).\"\"\"\n",
+    "        with torch.no_grad():\n",
+    "            for c in range(num_classes):\n",
+    "                mask = (labels == c)\n",
+    "                if mask.sum() < 2:\n",
+    "                    continue\n",
+    "                centroid = l2_normalize(embeddings[mask].mean(dim=0, keepdim=True))\n",
+    "                # Assign anchors to this class based on nearest centroid\n",
+    "                start = c * (self.num_anchors // num_classes)\n",
+    "                end = start + (self.num_anchors // num_classes)\n",
+    "                self.anchors.data[start:end] = l2_normalize(\n",
+    "                    (1 - momentum) * self.anchors.data[start:end] + momentum * centroid\n",
+    "                )\n",
+    "\n",
+    "    def anchor_activation(self, x):\n",
+    "        \"\"\"Return fraction of anchors that are 'active' (nearest to at least one sample).\"\"\"\n",
+    "        x_norm = l2_normalize(x)\n",
+    "        cos_sim = torch.mm(x_norm, self.anchors.T)\n",
+    "        nearest = cos_sim.argmax(dim=1)\n",
+    "        active = len(nearest.unique())\n",
+    "        return active, self.num_anchors\n",
+    "\n",
+    "# \u2500\u2500 Patchwork Block \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
+    "class SquaredReLU(nn.Module):\n",
+    "    def forward(self, x):\n",
+    "        return F.relu(x) ** 2\n",
+    "\n",
+    "class PatchworkBlock(nn.Module):\n",
+    "    \"\"\"SquaredReLU + LayerNorm block for post-constellation processing.\"\"\"\n",
+    "    def __init__(self, dim):\n",
+    "        super().__init__()\n",
+    "        self.net = nn.Sequential(\n",
+    "            nn.Linear(dim, dim),\n",
+    "            SquaredReLU(),\n",
+    "            nn.LayerNorm(dim),\n",
+    "            nn.Linear(dim, dim),\n",
+    "            SquaredReLU(),\n",
+    "            nn.LayerNorm(dim),\n",
+    "        )\n",
+    "    def forward(self, x):\n",
+    "        return x + self.net(x)  # residual\n",
+    "\n",
+    "class Patchwork(nn.Module):\n",
+    "    \"\"\"Stack of PatchworkBlocks.\"\"\"\n",
+    "    def __init__(self, dim, num_blocks=2):\n",
+    "        super().__init__()\n",
+    "        self.blocks = nn.Sequential(*[PatchworkBlock(dim) for _ in range(num_blocks)])\n",
+    "    def forward(self, x):\n",
+    "        return self.blocks(x)\n",
+    "\n",
+    "# \u2500\u2500 Classifier Head \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
+    "class ConstellationClassifier(nn.Module):\n",
+    "    \"\"\"Full pipeline: encoder output \u2192 constellation \u2192 patchwork \u2192 classify.\"\"\"\n",
+    "    def __init__(self, embed_dim, num_anchors=64, num_classes=10, patchwork_blocks=2):\n",
+    "        super().__init__()\n",
+    "        self.constellation = Constellation(embed_dim, num_anchors)\n",
+    "        self.patchwork = Patchwork(num_anchors, patchwork_blocks)\n",
+    "        self.head = nn.Linear(num_anchors, num_classes)\n",
+    "\n",
+    "    def forward(self, embeddings):\n",
+    "        tri = self.constellation(embeddings)       # (B, K)\n",
+    "        pw  = self.patchwork(tri)                   # (B, K)\n",
+    "        logits = self.head(pw)                      # (B, 10)\n",
+    "        return logits, tri\n",
+    "\n",
+    "    def count_params(self):\n",
+    "        return sum(p.numel() for p in self.parameters() if p.requires_grad)\n",
+    "\n",
+    "# \u2500\u2500 Losses \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
+    "class CayleyMengerVolumeLoss(nn.Module):\n",
+    "    \"\"\"CV loss: encourage constellation volume in 0.18-0.25 band.\"\"\"\n",
+    "    def __init__(self, target_low=0.18, target_high=0.25):\n",
+    "        super().__init__()\n",
+    "        self.target_low = target_low\n",
+    "        self.target_high = target_high\n",
+    "\n",
+    "    def forward(self, tri):\n",
+    "        \"\"\"tri: (B, K) triangulation distances.\"\"\"\n",
+    "        # Approximate volume via determinant of distance submatrix\n",
+    "        B, K = tri.shape\n",
+    "        # Use std of triangulation as proxy for volume spread\n",
+    "        vol = tri.std(dim=1).mean()\n",
+    "        # Penalty if outside target band\n",
+    "        below = F.relu(self.target_low - vol)\n",
+    "        above = F.relu(vol - self.target_high)\n",
+    "        return below + above, vol.item()\n",
+    "\n",
+    "class InfoNCELoss(nn.Module):\n",
+    "    \"\"\"Self-supervised contrastive loss on augmented views.\"\"\"\n",
+    "    def __init__(self, temperature=0.07):\n",
+    "        super().__init__()\n",
+    "        self.temperature = temperature\n",
+    "\n",
+    "    def forward(self, z1, z2):\n",
+    "        \"\"\"z1, z2: (B, d) L2-normalized embeddings of two augmented views.\"\"\"\n",
+    "        z1 = l2_normalize(z1)\n",
+    "        z2 = l2_normalize(z2)\n",
+    "        B = z1.shape[0]\n",
+    "        z = torch.cat([z1, z2], dim=0)  # (2B, d)\n",
+    "        sim = torch.mm(z, z.T) / self.temperature  # (2B, 2B)\n",
+    "        # Mask out self-similarity\n",
+    "        mask = ~torch.eye(2*B, dtype=bool, device=z.device)\n",
+    "        sim = sim.masked_fill(~mask, -1e9)\n",
+    "        # Positive pairs: (i, i+B) and (i+B, i)\n",
+    "        labels = torch.cat([torch.arange(B, 2*B), torch.arange(B)]).to(z.device)\n",
+    "        loss = F.cross_entropy(sim, labels)\n",
+    "        return loss\n",
+    "\n",
+    "class ClassConditionalInfoNCE(nn.Module):\n",
+    "    \"\"\"Class-conditional contrastive: positive = same class, not same image.\"\"\"\n",
+    "    def __init__(self, temperature=0.07):\n",
+    "        super().__init__()\n",
+    "        self.temperature = temperature\n",
+    "\n",
+    "    def forward(self, z, labels):\n",
+    "        z = l2_normalize(z)\n",
+    "        B = z.shape[0]\n",
+    "        sim = torch.mm(z, z.T) / self.temperature\n",
+    "        # Mask: positive if same class, negative otherwise\n",
+    "        label_mask = (labels.unsqueeze(0) == labels.unsqueeze(1)).float()\n",
+    "        label_mask.fill_diagonal_(0)  # exclude self\n",
+    "        # For each sample, compute log-softmax over all non-self, weighted by positives\n",
+    "        logits = sim - torch.diag(torch.full((B,), 1e9, device=z.device))\n",
+    "        pos_mask = label_mask\n",
+    "        neg_mask = 1 - label_mask\n",
+    "        neg_mask.fill_diagonal_(0)\n",
+    "        # InfoNCE: -log(sum_pos exp(sim) / sum_all exp(sim))\n",
+    "        exp_sim = torch.exp(logits)\n",
+    "        pos_sum = (exp_sim * pos_mask).sum(dim=1)\n",
+    "        all_sum = (exp_sim * (pos_mask + neg_mask)).sum(dim=1)\n",
+    "        loss = -torch.log(pos_sum / (all_sum + 1e-8) + 1e-8).mean()\n",
+    "        return loss\n",
+    "\n",
+    "# \u2500\u2500 TensorBoard Helpers \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
+    "def make_writer(experiment_id, base_dir=\"runs\"):\n",
+    "    \"\"\"Create a TensorBoard SummaryWriter with timestamped run directory.\"\"\"\n",
+    "    timestamp = datetime.datetime.now().strftime(\"%Y%m%d_%H%M%S\")\n",
+    "    log_dir = os.path.join(base_dir, f\"{experiment_id}_{timestamp}\")\n",
+    "    writer = SummaryWriter(log_dir=log_dir)\n",
+    "    print(f\"[TB] Logging to {log_dir}\")\n",
+    "    return writer\n",
+    "\n",
+    "def log_epoch_metrics(writer, epoch, metrics_dict, prefix=\"\"):\n",
+    "    \"\"\"Log a dict of metrics to TensorBoard.\"\"\"\n",
+    "    for key, val in metrics_dict.items():\n",
+    "        tag = f\"{prefix}/{key}\" if prefix else key\n",
+    "        writer.add_scalar(tag, val, epoch)\n",
+    "\n",
+    "def log_embedding_viz(writer, embeddings, labels, epoch, tag=\"embeddings\"):\n",
+    "    \"\"\"Log embeddings for TensorBoard projector.\"\"\"\n",
+    "    meta = [CIFAR10_CLASSES[l] for l in labels]\n",
+    "    writer.add_embedding(embeddings, metadata=meta, global_step=epoch, tag=tag)\n",
+    "\n",
+    "def log_anchor_histogram(writer, constellation, epoch):\n",
+    "    \"\"\"Log anchor norm distribution.\"\"\"\n",
+    "    norms = constellation.anchors.data.norm(dim=1)\n",
+    "    writer.add_histogram(\"anchors/norms\", norms, epoch)\n",
+    "    cos_sim_matrix = torch.mm(constellation.anchors.data, constellation.anchors.data.T)\n",
+    "    off_diag = cos_sim_matrix[~torch.eye(cos_sim_matrix.shape[0], dtype=bool)].flatten()\n",
+    "    writer.add_histogram(\"anchors/inter_cosine_sim\", off_diag, epoch)\n",
+    "\n",
+    "def log_confusion_matrix(writer, all_preds, all_labels, epoch):\n",
+    "    \"\"\"Log confusion matrix as image to TensorBoard.\"\"\"\n",
+    "    from sklearn.metrics import confusion_matrix\n",
+    "    cm = confusion_matrix(all_labels, all_preds, labels=range(10))\n",
+    "    fig, ax = plt.subplots(1, 1, figsize=(8, 8))\n",
+    "    im = ax.imshow(cm, interpolation=\"nearest\", cmap=plt.cm.Blues)\n",
+    "    ax.set_title(f\"Confusion Matrix \u2014 Epoch {epoch}\")\n",
+    "    ax.set_xlabel(\"Predicted\")\n",
+    "    ax.set_ylabel(\"True\")\n",
+    "    tick_marks = range(10)\n",
+    "    ax.set_xticks(tick_marks)\n",
+    "    ax.set_xticklabels(CIFAR10_CLASSES, rotation=45, ha=\"right\")\n",
+    "    ax.set_yticks(tick_marks)\n",
+    "    ax.set_yticklabels(CIFAR10_CLASSES)\n",
+    "    for i in range(10):\n",
+    "        for j in range(10):\n",
+    "            ax.text(j, i, str(cm[i, j]), ha=\"center\", va=\"center\",\n",
+    "                    color=\"white\" if cm[i, j] > cm.max() / 2 else \"black\", fontsize=7)\n",
+    "    fig.tight_layout()\n",
+    "    writer.add_figure(\"confusion_matrix\", fig, epoch)\n",
+    "    plt.close(fig)\n",
+    "\n",
+    "def log_triangulation_stats(writer, tri, epoch):\n",
+    "    \"\"\"Log triangulation distance statistics.\"\"\"\n",
+    "    writer.add_scalar(\"tri/mean\", tri.mean().item(), epoch)\n",
+    "    writer.add_scalar(\"tri/std\", tri.std().item(), epoch)\n",
+    "    writer.add_scalar(\"tri/min\", tri.min().item(), epoch)\n",
+    "    writer.add_scalar(\"tri/max\", tri.max().item(), epoch)\n",
+    "    writer.add_histogram(\"tri/distances\", tri.detach().cpu(), epoch)\n",
+    "\n",
+    "def log_scatter_2d(writer, embeddings, labels, epoch, tag=\"scatter_2d\", method=\"pca\"):\n",
+    "    \"\"\"Log 2D scatter plot of embeddings projected via PCA or t-SNE.\"\"\"\n",
+    "    emb_np = embeddings.detach().cpu().numpy()\n",
+    "    lab_np = labels.detach().cpu().numpy() if torch.is_tensor(labels) else np.array(labels)\n",
+    "    if method == \"pca\":\n",
+    "        from sklearn.decomposition import PCA\n",
+    "        proj = PCA(n_components=2).fit_transform(emb_np)\n",
+    "    else:\n",
+    "        from sklearn.manifold import TSNE\n",
+    "        proj = TSNE(n_components=2, perplexity=30, n_iter=500).fit_transform(emb_np)\n",
+    "    fig, ax = plt.subplots(1, 1, figsize=(8, 8))\n",
+    "    for c in range(10):\n",
+    "        mask = lab_np == c\n",
+    "        ax.scatter(proj[mask, 0], proj[mask, 1], s=3, alpha=0.5, label=CIFAR10_CLASSES[c])\n",
+    "    ax.legend(markerscale=5, fontsize=8)\n",
+    "    ax.set_title(f\"{tag} ({method.upper()}) \u2014 Epoch {epoch}\")\n",
+    "    fig.tight_layout()\n",
+    "    writer.add_figure(f\"{tag}/{method}\", fig, epoch)\n",
+    "    plt.close(fig)\n",
+    "\n",
+    "# \u2500\u2500 Training Utilities \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
+    "def train_one_epoch(model, loader, optimizer, criterion, device, epoch=0,\n",
+    "                    scheduler=None, extra_losses=None):\n",
+    "    \"\"\"\n",
+    "    Generic training loop for one epoch.\n",
+    "    model: callable(batch) \u2192 (logits, embeddings_on_sphere)\n",
+    "    extra_losses: list of (name, loss_fn, weight) where loss_fn(embeddings) \u2192 scalar\n",
+    "    \"\"\"\n",
+    "    model.train()\n",
+    "    total_loss = 0\n",
+    "    total_ce = 0\n",
+    "    correct = 0\n",
+    "    total = 0\n",
+    "    loss_components = defaultdict(float)\n",
+    "\n",
+    "    for batch_idx, (images, labels) in enumerate(loader):\n",
+    "        images, labels = images.to(device), labels.to(device)\n",
+    "        optimizer.zero_grad()\n",
+    "\n",
+    "        logits, embeddings = model(images)\n",
+    "        ce_loss = criterion(logits, labels)\n",
+    "        loss = ce_loss\n",
+    "        loss_components[\"CE\"] += ce_loss.item()\n",
+    "\n",
+    "        if extra_losses:\n",
+    "            for name, loss_fn, weight in extra_losses:\n",
+    "                if name == \"CV\":\n",
+    "                    el, vol_val = loss_fn(embeddings)\n",
+    "                    loss_components[\"CV_vol\"] += vol_val\n",
+    "                elif name == \"InfoNCE_class\":\n",
+    "                    el = loss_fn(embeddings, labels)\n",
+    "                else:\n",
+    "                    el = loss_fn(embeddings)\n",
+    "                loss = loss + weight * el\n",
+    "                loss_components[name] += el.item()\n",
+    "\n",
+    "        loss.backward()\n",
+    "        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)\n",
+    "        optimizer.step()\n",
+    "        if scheduler:\n",
+    "            scheduler.step()\n",
+    "\n",
+    "        total_loss += loss.item()\n",
+    "        _, predicted = logits.max(1)\n",
+    "        total += labels.size(0)\n",
+    "        correct += predicted.eq(labels).sum().item()\n",
+    "\n",
+    "    n_batches = len(loader)\n",
+    "    metrics = {\n",
+    "        \"loss/total\": total_loss / n_batches,\n",
+    "        \"loss/CE\": loss_components[\"CE\"] / n_batches,\n",
+    "        \"accuracy/train\": 100.0 * correct / total,\n",
+    "    }\n",
+    "    for k, v in loss_components.items():\n",
+    "        if k not in (\"CE\",):\n",
+    "            metrics[f\"loss/{k}\"] = v / n_batches\n",
+    "    return metrics\n",
+    "\n",
+    "@torch.no_grad()\n",
+    "def evaluate(model, loader, criterion, device, collect_embeddings=False):\n",
+    "    \"\"\"Evaluate model on validation set. Optionally collect embeddings.\"\"\"\n",
+    "    model.eval()\n",
+    "    total_loss = 0\n",
+    "    correct = 0\n",
+    "    total = 0\n",
+    "    all_preds = []\n",
+    "    all_labels = []\n",
+    "    all_embeddings = []\n",
+    "\n",
+    "    for images, labels in loader:\n",
+    "        images, labels = images.to(device), labels.to(device)\n",
+    "        logits, embeddings = model(images)\n",
+    "        loss = criterion(logits, labels)\n",
+    "        total_loss += loss.item()\n",
+    "        _, predicted = logits.max(1)\n",
+    "        total += labels.size(0)\n",
+    "        correct += predicted.eq(labels).sum().item()\n",
+    "        all_preds.extend(predicted.cpu().numpy())\n",
+    "        all_labels.extend(labels.cpu().numpy())\n",
+    "        if collect_embeddings:\n",
+    "            all_embeddings.append(embeddings.cpu())\n",
+    "\n",
+    "    n_batches = len(loader)\n",
+    "    metrics = {\n",
+    "        \"loss/val\": total_loss / n_batches,\n",
+    "        \"accuracy/val\": 100.0 * correct / total,\n",
+    "    }\n",
+    "    result = {\"metrics\": metrics, \"preds\": np.array(all_preds), \"labels\": np.array(all_labels)}\n",
+    "    if collect_embeddings:\n",
+    "        result[\"embeddings\"] = torch.cat(all_embeddings, dim=0)\n",
+    "    return result\n",
+    "\n",
+    "def run_experiment(experiment_id, model, train_loader, val_loader, device,\n",
+    "                   epochs=100, lr=1e-3, weight_decay=1e-4, extra_losses=None,\n",
+    "                   push_anchors_every=50, log_embeddings_every=25,\n",
+    "                   patience=20):\n",
+    "    \"\"\"\n",
+    "    Full experiment runner with TensorBoard logging, early stopping, and diagnostics.\n",
+    "    model must return (logits, sphere_embeddings).\n",
+    "    \"\"\"\n",
+    "    writer = make_writer(experiment_id)\n",
+    "    criterion = nn.CrossEntropyLoss()\n",
+    "    optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=weight_decay)\n",
+    "    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs)\n",
+    "\n",
+    "    best_val_acc = 0.0\n",
+    "    best_epoch = 0\n",
+    "    no_improve = 0\n",
+    "\n",
+    "    # Log model info\n",
+    "    n_params = sum(p.numel() for p in model.parameters() if p.requires_grad)\n",
+    "    writer.add_text(\"config/params\", f\"Learned params: {n_params:,}\")\n",
+    "    writer.add_text(\"config/experiment\", experiment_id)\n",
+    "    print(f\"\\n{'='*70}\")\n",
+    "    print(f\"[EXP] {experiment_id}\")\n",
+    "    print(f\"[EXP] Learned params: {n_params:,}, Epochs: {epochs}, LR: {lr}\")\n",
+    "    print(f\"{'='*70}\")\n",
+    "\n",
+    "    for epoch in range(1, epochs + 1):\n",
+    "        t0 = time.time()\n",
+    "\n",
+    "        # Train\n",
+    "        train_metrics = train_one_epoch(\n",
+    "            model, train_loader, optimizer, criterion, device, epoch,\n",
+    "            extra_losses=extra_losses\n",
+    "        )\n",
+    "\n",
+    "        # Evaluate\n",
+    "        collect_emb = (epoch % log_embeddings_every == 0) or (epoch == epochs)\n",
+    "        val_result = evaluate(model, val_loader, criterion, device, collect_embeddings=collect_emb)\n",
+    "        val_metrics = val_result[\"metrics\"]\n",
+    "\n",
+    "        # Combine metrics\n",
+    "        all_metrics = {**train_metrics, **val_metrics}\n",
+    "        all_metrics[\"lr\"] = optimizer.param_groups[0][\"lr\"]\n",
+    "        all_metrics[\"epoch_time\"] = time.time() - t0\n",
+    "        all_metrics[\"train_val_gap\"] = train_metrics[\"accuracy/train\"] - val_metrics[\"accuracy/val\"]\n",
+    "\n",
+    "        # Constellation diagnostics\n",
+    "        if hasattr(model, \"constellation\"):\n",
+    "            constellation = model.constellation\n",
+    "        elif hasattr(model, \"classifier\") and hasattr(model.classifier, \"constellation\"):\n",
+    "            constellation = model.classifier.constellation\n",
+    "        else:\n",
+    "            constellation = None\n",
+    "\n",
+    "        if constellation is not None:\n",
+    "            log_anchor_histogram(writer, constellation, epoch)\n",
+    "\n",
+    "            # Push anchors periodically\n",
+    "            if epoch % push_anchors_every == 0 and \"embeddings\" in val_result:\n",
+    "                emb = val_result[\"embeddings\"][:2000].to(device)\n",
+    "                lab = torch.tensor(val_result[\"labels\"][:2000]).to(device)\n",
+    "                constellation.push_anchors(emb, lab)\n",
+    "                active, total_anch = constellation.anchor_activation(emb)\n",
+    "                all_metrics[\"anchors/active\"] = active\n",
+    "                all_metrics[\"anchors/total\"] = total_anch\n",
+    "                all_metrics[\"anchors/utilization\"] = active / total_anch\n",
+    "                print(f\"  [ANCHOR PUSH] Active: {active}/{total_anch} ({100*active/total_anch:.1f}%)\")\n",
+    "\n",
+    "        # Log all metrics\n",
+    "        log_epoch_metrics(writer, epoch, all_metrics)\n",
+    "        scheduler.step()\n",
+    "\n",
+    "        # Log embeddings visualization\n",
+    "        if collect_emb and \"embeddings\" in val_result:\n",
+    "            emb = val_result[\"embeddings\"][:2000]\n",
+    "            lab = val_result[\"labels\"][:2000]\n",
+    "            log_scatter_2d(writer, emb, lab, epoch, tag=\"val_embeddings\", method=\"pca\")\n",
+    "            if epoch == epochs:\n",
+    "                log_confusion_matrix(writer, val_result[\"preds\"], val_result[\"labels\"], epoch)\n",
+    "                log_scatter_2d(writer, emb, lab, epoch, tag=\"val_embeddings_tsne\", method=\"tsne\")\n",
+    "\n",
+    "        # Log triangulation stats from a sample batch\n",
+    "        if constellation is not None and collect_emb and \"embeddings\" in val_result:\n",
+    "            sample_emb = val_result[\"embeddings\"][:500].to(device)\n",
+    "            tri_sample = constellation(sample_emb)\n",
+    "            log_triangulation_stats(writer, tri_sample, epoch)\n",
+    "\n",
+    "        # Best model tracking\n",
+    "        val_acc = val_metrics[\"accuracy/val\"]\n",
+    "        if val_acc > best_val_acc:\n",
+    "            best_val_acc = val_acc\n",
+    "            best_epoch = epoch\n",
+    "            no_improve = 0\n",
+    "            writer.add_scalar(\"best/val_acc\", best_val_acc, epoch)\n",
+    "            writer.add_scalar(\"best/epoch\", best_epoch, epoch)\n",
+    "        else:\n",
+    "            no_improve += 1\n",
+    "\n",
+    "        # Print progress\n",
+    "        if epoch % 5 == 0 or epoch == 1 or epoch == epochs:\n",
+    "            print(f\"  E{epoch:3d} | TrainAcc {train_metrics['accuracy/train']:5.1f}% | \"\n",
+    "                  f\"ValAcc {val_acc:5.1f}% | Loss {all_metrics['loss/total']:.4f} | \"\n",
+    "                  f\"Best {best_val_acc:.1f}% @E{best_epoch} | \"\n",
+    "                  f\"Gap {all_metrics['train_val_gap']:+.1f}% | \"\n",
+    "                  f\"{all_metrics['epoch_time']:.1f}s\")\n",
+    "\n",
+    "        # Early stopping\n",
+    "        if no_improve >= patience:\n",
+    "            print(f\"  [EARLY STOP] No improvement for {patience} epochs. Best: {best_val_acc:.1f}% @E{best_epoch}\")\n",
+    "            break\n",
+    "\n",
+    "    writer.close()\n",
+    "    print(f\"\\n[RESULT] {experiment_id}: Best Val Acc = {best_val_acc:.2f}% at epoch {best_epoch}\")\n",
+    "    print(f\"[RESULT] Learned params: {n_params:,}\")\n",
+    "    return {\"experiment\": experiment_id, \"best_val_acc\": best_val_acc, \"best_epoch\": best_epoch,\n",
+    "            \"params\": n_params, \"final_train_acc\": train_metrics[\"accuracy/train\"]}\n",
+    "\n",
+    "# \u2500\u2500 Johnson-Lindenstrauss Projection \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
+    "class JLProjection(nn.Module):\n",
+    "    \"\"\"Fixed random Gaussian projection (not learned).\"\"\"\n",
+    "    def __init__(self, input_dim, output_dim):\n",
+    "        super().__init__()\n",
+    "        self.register_buffer(\"W\", torch.randn(input_dim, output_dim) / math.sqrt(output_dim))\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        return x @ self.W\n",
+    "\n",
+    "# \u2500\u2500 Results Scoreboard \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
+    "class Scoreboard:\n",
+    "    \"\"\"Track results across experiments in a session.\"\"\"\n",
+    "    def __init__(self):\n",
+    "        self.results = []\n",
+    "\n",
+    "    def add(self, result_dict):\n",
+    "        self.results.append(result_dict)\n",
+    "        self._print_table()\n",
+    "\n",
+    "    def _print_table(self):\n",
+    "        print(f\"\\n{'='*80}\")\n",
+    "        print(f\"{'SCOREBOARD':^80}\")\n",
+    "        print(f\"{'='*80}\")\n",
+    "        print(f\"{'Experiment':<40} {'Val Acc':>8} {'Params':>10} {'Best E':>7}\")\n",
+    "        print(f\"{'-'*40} {'-'*8} {'-'*10} {'-'*7}\")\n",
+    "        for r in sorted(self.results, key=lambda x: x[\"best_val_acc\"], reverse=True):\n",
+    "            print(f\"{r['experiment']:<40} {r['best_val_acc']:>7.2f}% {r['params']:>10,} {r['best_epoch']:>7}\")\n",
+    "        print(f\"{'='*80}\\n\")\n",
+    "\n",
+    "    def save(self, path=\"scoreboard.json\"):\n",
+    "        with open(path, \"w\") as f:\n",
+    "            json.dump(self.results, f, indent=2)\n",
+    "        print(f\"[SCOREBOARD] Saved to {path}\")\n",
+    "\n",
+    "scoreboard = Scoreboard()\n",
+    "print(\"[SHARED] All utilities loaded successfully.\")\n",
+    ""
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Additional Pipeline Components"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "\n",
+    "from kymatio.torch import Scattering2D\n",
+    "import scipy.signal as signal\n",
+    "\n",
+    "class CurveletEnergy(nn.Module):\n",
+    "    '''Approximate curvelet energy (simplified ridge detection).'''\n",
+    "    def __init__(self):\n",
+    "        super().__init__()\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        # Simplified: use Sobel edge detection as proxy for curvelet energy\n",
+    "        B, C, H, W = x.shape\n",
+    "        energy = []\n",
+    "\n",
+    "        for c in range(C):\n",
+    "            channel = x[:, c:c+1, :, :]\n",
+    "            # Compute gradients\n",
+    "            dx = torch.abs(channel[:, :, 1:, :] - channel[:, :, :-1, :])\n",
+    "            dy = torch.abs(channel[:, :, :, 1:] - channel[:, :, :, :-1:])\n",
+    "            # Energy as magnitude\n",
+    "            e = (dx.mean(dim=(2, 3)) + dy.mean(dim=(2, 3)))\n",
+    "            energy.append(e)\n",
+    "\n",
+    "        return torch.cat(energy, dim=1)  # (B, 3)\n",
+    "\n",
+    "class RadonTransform(nn.Module):\n",
+    "    '''Approximate Radon transform via rotated projections.'''\n",
+    "    def __init__(self, n_angles=16):\n",
+    "        super().__init__()\n",
+    "        self.n_angles = n_angles\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        # Simplified: project image at different angles\n",
+    "        B, C, H, W = x.shape\n",
+    "        radon_feats = []\n",
+    "\n",
+    "        for c in range(C):\n",
+    "            channel = x[:, c:c+1, :, :]  # (B, 1, H, W)\n",
+    "            angles = []\n",
+    "            for theta in torch.linspace(0, np.pi, self.n_angles):\n",
+    "                # Project along angle (simplified)\n",
+    "                proj = channel.mean(dim=3 if abs(np.cos(theta)) > abs(np.sin(theta)) else 2)\n",
+    "                angles.append(proj.mean(dim=2, keepdim=True))\n",
+    "            radon_feats.extend(angles)\n",
+    "\n",
+    "        return torch.cat(radon_feats, dim=1)  # (B, 3*n_angles)\n",
+    "\n",
+    "print(\"[PIPELINES] Additional components loaded.\")\n",
+    ""
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 7.1a: Reference Pipeline\n",
+    "\n",
+    "Scattering(J=2,L=8) \u2192 JL(128) \u2192 L2 norm \u2192 constellation(64) \u2192 classify"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "\n",
+    "print(\"[7.1a] Reference Pipeline\")\n",
+    "print(\"=\"*70)\n",
+    "\n",
+    "class ReferencePipeline(nn.Module):\n",
+    "    '''Canonical: Scattering \u2192 JL \u2192 constellation \u2192 classify.'''\n",
+    "    def __init__(self, embed_dim=128, num_classes=10):\n",
+    "        super().__init__()\n",
+    "        self.scattering = Scattering2D(J=2, L=8, shape=(32, 32), max_order=2)\n",
+    "        # Scattering output: 81 channels\n",
+    "        self.jl = JLProjection(81, embed_dim)\n",
+    "        self.classifier = ConstellationClassifier(embed_dim, num_anchors=64, num_classes=num_classes)\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        scatt = self.scattering(x)  # (B, 81, 32, 32)\n",
+    "        # Pool to 1D\n",
+    "        scatt_pool = F.adaptive_avg_pool2d(scatt, (1, 1)).squeeze(-1).squeeze(-1)  # (B, 81)\n",
+    "        emb = l2_normalize(self.jl(scatt_pool))  # (B, 128)\n",
+    "        logits, tri = self.classifier(emb)\n",
+    "        return logits, emb\n",
+    "\n",
+    "model = ReferencePipeline(embed_dim=128, num_classes=10).to(device)\n",
+    "print(f\"[7.1a] Model params: {sum(p.numel() for p in model.parameters()):,}\")\n",
+    "\n",
+    "train_loader, val_loader = get_cifar10_loaders(batch_size=256, augment=True)\n",
+    "\n",
+    "result_7_1a = run_experiment(\n",
+    "    \"7.1a_Reference_Pipeline\",\n",
+    "    model, train_loader, val_loader, device,\n",
+    "    epochs=100, lr=1e-3, weight_decay=1e-4, patience=20\n",
+    ")\n",
+    "scoreboard.add(result_7_1a)\n",
+    ""
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 7.1b: Learned 2-Layer Projection\n",
+    "\n",
+    "Replace JL with learned linear layer (~16K params)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "\n",
+    "print(\"[7.1b] Learned 2-Layer Projection\")\n",
+    "print(\"=\"*70)\n",
+    "\n",
+    "class LearnedProjectionPipeline(nn.Module):\n",
+    "    '''Scattering \u2192 learned 2-layer projection \u2192 constellation.'''\n",
+    "    def __init__(self, embed_dim=128, num_classes=10, hidden_dim=128):\n",
+    "        super().__init__()\n",
+    "        self.scattering = Scattering2D(J=2, L=8, shape=(32, 32), max_order=2)\n",
+    "        # Learned projection\n",
+    "        self.proj = nn.Sequential(\n",
+    "            nn.Linear(81, hidden_dim),\n",
+    "            nn.ReLU(),\n",
+    "            nn.Linear(hidden_dim, embed_dim)\n",
+    "        )\n",
+    "        self.classifier = ConstellationClassifier(embed_dim, num_anchors=64, num_classes=num_classes)\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        scatt = self.scattering(x)\n",
+    "        scatt_pool = F.adaptive_avg_pool2d(scatt, (1, 1)).squeeze(-1).squeeze(-1)\n",
+    "        emb = l2_normalize(self.proj(scatt_pool))\n",
+    "        logits, tri = self.classifier(emb)\n",
+    "        return logits, emb\n",
+    "\n",
+    "model = LearnedProjectionPipeline(embed_dim=128, num_classes=10, hidden_dim=128).to(device)\n",
+    "params = sum(p.numel() for p in model.parameters())\n",
+    "print(f\"[7.1b] Model params: {params:,}\")\n",
+    "\n",
+    "result_7_1b = run_experiment(\n",
+    "    \"7.1b_Learned_Projection\",\n",
+    "    model, train_loader, val_loader, device,\n",
+    "    epochs=100, lr=1e-3, weight_decay=1e-4, patience=20\n",
+    ")\n",
+    "scoreboard.add(result_7_1b)\n",
+    ""
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 7.1c: Scattering + Curvelet Energy\n",
+    "\n",
+    "Scattering \u2192 curvelet energy \u2192 concat \u2192 JL"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "\n",
+    "print(\"[7.1c] Scattering + Curvelet Energy\")\n",
+    "print(\"=\"*70)\n",
+    "\n",
+    "class ScatteringCurveletPipeline(nn.Module):\n",
+    "    '''Scattering + curvelet energy concat.'''\n",
+    "    def __init__(self, embed_dim=128, num_classes=10):\n",
+    "        super().__init__()\n",
+    "        self.scattering = Scattering2D(J=2, L=8, shape=(32, 32), max_order=2)\n",
+    "        self.curvelet = CurveletEnergy()\n",
+    "        # Scattering (81) + curvelet (3) = 84\n",
+    "        self.jl = JLProjection(84, embed_dim)\n",
+    "        self.classifier = ConstellationClassifier(embed_dim, num_anchors=64, num_classes=num_classes)\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        scatt = self.scattering(x)\n",
+    "        scatt_pool = F.adaptive_avg_pool2d(scatt, (1, 1)).squeeze(-1).squeeze(-1)\n",
+    "        curv = self.curvelet(x)\n",
+    "        combined = torch.cat([scatt_pool, curv], dim=1)\n",
+    "        emb = l2_normalize(self.jl(combined))\n",
+    "        logits, tri = self.classifier(emb)\n",
+    "        return logits, emb\n",
+    "\n",
+    "model = ScatteringCurveletPipeline(embed_dim=128, num_classes=10).to(device)\n",
+    "print(f\"[7.1c] Model params: {sum(p.numel() for p in model.parameters()):,}\")\n",
+    "\n",
+    "result_7_1c = run_experiment(\n",
+    "    \"7.1c_Scattering_Curvelet\",\n",
+    "    model, train_loader, val_loader, device,\n",
+    "    epochs=100, lr=1e-3, weight_decay=1e-4, patience=20\n",
+    ")\n",
+    "scoreboard.add(result_7_1c)\n",
+    ""
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 7.2a: STFT Channels + Scattering\n",
+    "\n",
+    "STFT channels (v4) + scattering \u2192 concat \u2192 JL \u2192 constellation"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "\n",
+    "print(\"[7.2a] STFT Channels + Scattering\")\n",
+    "print(\"=\"*70)\n",
+    "\n",
+    "class STFTScatteringPipeline(nn.Module):\n",
+    "    '''STFT + Scattering combined features.'''\n",
+    "    def __init__(self, embed_dim=128, num_classes=10):\n",
+    "        super().__init__()\n",
+    "        self.scattering = Scattering2D(J=2, L=8, shape=(32, 32), max_order=2)\n",
+    "        # STFT front-end (from nb5)\n",
+    "        self.stft_fe = STFTFrontend(n_fft=16, hop_length=4, freq_bins=8)\n",
+    "        # STFT (1536) + Scattering (81) = 1617, reduce to ~256\n",
+    "        self.reduce = nn.Linear(1617, 256)\n",
+    "        self.jl = JLProjection(256, embed_dim)\n",
+    "        self.classifier = ConstellationClassifier(embed_dim, num_anchors=64, num_classes=num_classes)\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        # Scattering\n",
+    "        scatt = self.scattering(x)\n",
+    "        scatt_pool = F.adaptive_avg_pool2d(scatt, (1, 1)).squeeze(-1).squeeze(-1)\n",
+    "\n",
+    "        # STFT\n",
+    "        stft_feat = self.stft_fe(x)\n",
+    "\n",
+    "        # Combine\n",
+    "        combined = torch.cat([scatt_pool, stft_feat], dim=1)\n",
+    "        reduced = self.reduce(combined)\n",
+    "        emb = l2_normalize(self.jl(reduced))\n",
+    "        logits, tri = self.classifier(emb)\n",
+    "        return logits, emb\n",
+    "\n",
+    "model = STFTScatteringPipeline(embed_dim=128, num_classes=10).to(device)\n",
+    "print(f\"[7.2a] Model params: {sum(p.numel() for p in model.parameters()):,}\")\n",
+    "\n",
+    "result_7_2a = run_experiment(\n",
+    "    \"7.2a_STFT_Scattering\",\n",
+    "    model, train_loader, val_loader, device,\n",
+    "    epochs=100, lr=1e-3, weight_decay=1e-4, patience=20\n",
+    ")\n",
+    "scoreboard.add(result_7_2a)\n",
+    ""
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 7.2b: Scattering \u2192 Cholesky Gram + SVD Signature\n",
+    "\n",
+    "Scattering \u2192 Cholesky of Gram + SVD singular values"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "\n",
+    "print(\"[7.2b] Scattering + Cholesky Gram + SVD\")\n",
+    "print(\"=\"*70)\n",
+    "\n",
+    "class ScatteringCholeskyPipeline(nn.Module):\n",
+    "    '''Scattering \u2192 Gram \u2192 Cholesky + SVD signatures.'''\n",
+    "    def __init__(self, embed_dim=128, num_classes=10):\n",
+    "        super().__init__()\n",
+    "        self.scattering = Scattering2D(J=2, L=8, shape=(32, 32), max_order=2)\n",
+    "        # Scattering (81) \u2192 reduced (256) for Gram\n",
+    "        self.reduce = nn.Linear(81, 256)\n",
+    "\n",
+    "        # Features: Cholesky (10) + SVD (10) + original (81)\n",
+    "        self.jl = JLProjection(101, embed_dim)\n",
+    "        self.classifier = ConstellationClassifier(embed_dim, num_anchors=64, num_classes=num_classes)\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        scatt = self.scattering(x)\n",
+    "        scatt_pool = F.adaptive_avg_pool2d(scatt, (1, 1)).squeeze(-1).squeeze(-1)\n",
+    "\n",
+    "        # Gram matrix\n",
+    "        feat_red = self.reduce(scatt_pool)\n",
+    "        B = feat_red.shape[0]\n",
+    "        gram = torch.bmm(feat_red.unsqueeze(2), feat_red.unsqueeze(1))\n",
+    "        gram = gram + torch.eye(256, device=gram.device).unsqueeze(0) * 1e-6\n",
+    "\n",
+    "        # Cholesky\n",
+    "        try:\n",
+    "            L = torch.linalg.cholesky(gram)\n",
+    "            chol_diag = torch.diagonal(L, dim1=1, dim2=2)\n",
+    "            chol_mean = chol_diag.mean(dim=1)\n",
+    "            chol_std = chol_diag.std(dim=1)\n",
+    "            chol_sum = chol_diag.sum(dim=1)\n",
+    "            chol_max = chol_diag.max(dim=1)[0]\n",
+    "            chol_norm = torch.linalg.norm(chol_diag, dim=1)\n",
+    "            chol_features = torch.stack([chol_mean, chol_std, chol_sum, chol_max, chol_norm,\n",
+    "                                         chol_diag.min(dim=1)[0], chol_diag.max(dim=1)[0] - chol_diag.min(dim=1)[0],\n",
+    "                                         torch.prod(chol_diag, dim=1), torch.log(torch.prod(chol_diag, dim=1) + 1e-8),\n",
+    "                                         (chol_diag > 1e-5).sum(dim=1).float()], dim=1)\n",
+    "        except:\n",
+    "            chol_features = torch.zeros(B, 10, device=x.device)\n",
+    "\n",
+    "        # SVD\n",
+    "        U, S, Vh = torch.svd(gram)\n",
+    "        svd_mean = S.mean(dim=1)\n",
+    "        svd_std = S.std(dim=1)\n",
+    "        svd_ratio = S[:, 0] / (S[:, -1] + 1e-8)\n",
+    "        svd_sum = S.sum(dim=1)\n",
+    "        svd_max = S[:, 0]\n",
+    "        svd_min = S[:, -1]\n",
+    "        svd_norm = torch.linalg.norm(S, dim=1)\n",
+    "        svd_rank = (S > 1e-5).sum(dim=1).float()\n",
+    "        svd_trace = S.sum(dim=1)\n",
+    "        svd_energy = (S ** 2).sum(dim=1)\n",
+    "        svd_features = torch.stack([svd_mean, svd_std, svd_ratio, svd_sum, svd_max,\n",
+    "                                    svd_min, svd_norm, svd_rank, svd_trace, svd_energy], dim=1)\n",
+    "\n",
+    "        # Combine\n",
+    "        combined = torch.cat([scatt_pool, chol_features, svd_features], dim=1)  # (B, 101)\n",
+    "        emb = l2_normalize(self.jl(combined))\n",
+    "        logits, tri = self.classifier(emb)\n",
+    "        return logits, emb\n",
+    "\n",
+    "model = ScatteringCholeskyPipeline(embed_dim=128, num_classes=10).to(device)\n",
+    "print(f\"[7.2b] Model params: {sum(p.numel() for p in model.parameters()):,}\")\n",
+    "\n",
+    "result_7_2b = run_experiment(\n",
+    "    \"7.2b_Scattering_Cholesky_SVD\",\n",
+    "    model, train_loader, val_loader, device,\n",
+    "    epochs=100, lr=1e-3, weight_decay=1e-4, patience=20\n",
+    ")\n",
+    "scoreboard.add(result_7_2b)\n",
+    ""
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 7.3a: Multi-Signature Parallel\n",
+    "\n",
+    "Scattering + Gabor + Radon \u2192 separate constellations \u2192 fusion"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "\n",
+    "print(\"[7.3a] Multi-Signature Parallel\")\n",
+    "print(\"=\"*70)\n",
+    "\n",
+    "class GaborFilter(nn.Module):\n",
+    "    '''Simple Gabor filter bank approximation.'''\n",
+    "    def __init__(self, n_orientations=8, n_scales=3):\n",
+    "        super().__init__()\n",
+    "        self.n_orientations = n_orientations\n",
+    "        self.n_scales = n_scales\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        B, C, H, W = x.shape\n",
+    "        outputs = []\n",
+    "\n",
+    "        for scale in range(self.n_scales):\n",
+    "            sigma = 1.0 + scale * 0.5\n",
+    "            lambda_val = 5.0 + scale\n",
+    "            for theta in torch.linspace(0, np.pi, self.n_orientations):\n",
+    "                # Simplified Gabor: Gaussian-weighted sinusoid\n",
+    "                y = torch.arange(H, dtype=x.dtype, device=x.device)\n",
+    "                x_grid = torch.arange(W, dtype=x.dtype, device=x.device)\n",
+    "                yy, xx = torch.meshgrid(y - H/2, x_grid - W/2, indexing='ij')\n",
+    "\n",
+    "                # Rotate\n",
+    "                x_theta = xx * np.cos(theta) + yy * np.sin(theta)\n",
+    "                y_theta = -xx * np.sin(theta) + yy * np.cos(theta)\n",
+    "\n",
+    "                # Gabor kernel\n",
+    "                kernel = torch.exp(-(x_theta**2 + y_theta**2) / (2 * sigma**2)) *                          torch.cos(2 * np.pi * x_theta / lambda_val)\n",
+    "                kernel = kernel / (kernel.abs().sum() + 1e-8)\n",
+    "\n",
+    "                # Convolve\n",
+    "                kernel = kernel.unsqueeze(0).unsqueeze(0)\n",
+    "                response = F.conv2d(x, kernel.expand(C, -1, -1, -1), padding='same').mean(dim=(2, 3))\n",
+    "                outputs.append(response)\n",
+    "\n",
+    "        return torch.cat(outputs, dim=1)\n",
+    "\n",
+    "class MultiSignaturePipeline(nn.Module):\n",
+    "    '''Parallel: Scattering, Gabor, Radon \u2192 fusion.'''\n",
+    "    def __init__(self, embed_dim=128, num_classes=10):\n",
+    "        super().__init__()\n",
+    "        self.scattering = Scattering2D(J=2, L=8, shape=(32, 32), max_order=2)\n",
+    "        self.gabor = GaborFilter(n_orientations=8, n_scales=2)\n",
+    "        self.radon = RadonTransform(n_angles=8)\n",
+    "\n",
+    "        # Create separate branches\n",
+    "        self.scatt_jl = JLProjection(81, 64)\n",
+    "        self.gabor_jl = JLProjection(48, 64)  # 8*2 orientations/scales * 3 channels\n",
+    "        self.radon_jl = JLProjection(24, 64)  # 8 angles * 3 channels\n",
+    "\n",
+    "        # Fusion\n",
+    "        self.fusion = nn.Linear(192, embed_dim)\n",
+    "\n",
+    "        # Classifier (3 constellations or single)\n",
+    "        self.classifier = ConstellationClassifier(embed_dim, num_anchors=64, num_classes=num_classes)\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        # Scattering\n",
+    "        scatt = self.scattering(x)\n",
+    "        scatt_pool = F.adaptive_avg_pool2d(scatt, (1, 1)).squeeze(-1).squeeze(-1)\n",
+    "        scatt_emb = l2_normalize(self.scatt_jl(scatt_pool))\n",
+    "\n",
+    "        # Gabor\n",
+    "        gabor = self.gabor(x)\n",
+    "        gabor_emb = l2_normalize(self.gabor_jl(gabor))\n",
+    "\n",
+    "        # Radon\n",
+    "        radon = self.radon(x)\n",
+    "        radon_emb = l2_normalize(self.radon_jl(radon))\n",
+    "\n",
+    "        # Fuse\n",
+    "        fused = torch.cat([scatt_emb, gabor_emb, radon_emb], dim=1)\n",
+    "        emb = l2_normalize(self.fusion(fused))\n",
+    "\n",
+    "        logits, tri = self.classifier(emb)\n",
+    "        return logits, emb\n",
+    "\n",
+    "model = MultiSignaturePipeline(embed_dim=128, num_classes=10).to(device)\n",
+    "print(f\"[7.3a] Model params: {sum(p.numel() for p in model.parameters()):,}\")\n",
+    "\n",
+    "result_7_3a = run_experiment(\n",
+    "    \"7.3a_MultiSignature_Parallel\",\n",
+    "    model, train_loader, val_loader, device,\n",
+    "    epochs=100, lr=1e-3, weight_decay=1e-4, patience=20\n",
+    ")\n",
+    "scoreboard.add(result_7_3a)\n",
+    ""
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 7.3b: Hierarchical Constellation\n",
+    "\n",
+    "Scattering \u2192 coarse anchors \u2192 residual \u2192 fine anchors"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "\n",
+    "print(\"[7.3b] Hierarchical Constellation\")\n",
+    "print(\"=\"*70)\n",
+    "\n",
+    "class HierarchicalConstellation(nn.Module):\n",
+    "    '''Two-level constellation: coarse + residual fine.'''\n",
+    "    def __init__(self, embed_dim, num_coarse=16, num_fine=48):\n",
+    "        super().__init__()\n",
+    "        self.embed_dim = embed_dim\n",
+    "        self.num_coarse = num_coarse\n",
+    "        self.num_fine = num_fine\n",
+    "\n",
+    "        # Coarse constellation\n",
+    "        coarse_anchors = self._init_anchors(embed_dim, num_coarse, \"uniform\")\n",
+    "        self.coarse_anchors = nn.Parameter(coarse_anchors, requires_grad=False)\n",
+    "\n",
+    "        # Fine constellation\n",
+    "        fine_anchors = self._init_anchors(embed_dim, num_fine, \"uniform\")\n",
+    "        self.fine_anchors = nn.Parameter(fine_anchors, requires_grad=False)\n",
+    "\n",
+    "    def _init_anchors(self, d, K, method):\n",
+    "        A = torch.randn(K, d)\n",
+    "        Q, _ = torch.linalg.qr(A.T)\n",
+    "        anchors = Q.T[:K]\n",
+    "        if K > d:\n",
+    "            extra = F.normalize(torch.randn(K - d, d), dim=-1)\n",
+    "            anchors = torch.cat([anchors, extra], dim=0)\n",
+    "        return F.normalize(anchors, dim=-1)\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        x_norm = l2_normalize(x)\n",
+    "\n",
+    "        # Coarse triangulation\n",
+    "        coarse_sim = torch.mm(x_norm, self.coarse_anchors.T)\n",
+    "        coarse_tri = 1.0 - coarse_sim  # (B, 16)\n",
+    "\n",
+    "        # Residual toward nearest coarse anchor\n",
+    "        nearest_coarse = coarse_sim.argmax(dim=1)\n",
+    "        nearest_anchor = self.coarse_anchors[nearest_coarse]  # (B, d)\n",
+    "        residual = x_norm - nearest_anchor  # (B, d)\n",
+    "        residual = l2_normalize(residual)\n",
+    "\n",
+    "        # Fine triangulation (relative to residual)\n",
+    "        fine_sim = torch.mm(residual, self.fine_anchors.T)\n",
+    "        fine_tri = 1.0 - fine_sim  # (B, 48)\n",
+    "\n",
+    "        # Combine\n",
+    "        tri = torch.cat([coarse_tri, fine_tri], dim=1)  # (B, 64)\n",
+    "        return tri\n",
+    "\n",
+    "class HierarchicalPipeline(nn.Module):\n",
+    "    def __init__(self, embed_dim=128, num_classes=10):\n",
+    "        super().__init__()\n",
+    "        self.scattering = Scattering2D(J=2, L=8, shape=(32, 32), max_order=2)\n",
+    "        self.jl = JLProjection(81, embed_dim)\n",
+    "        self.constellation = HierarchicalConstellation(embed_dim, num_coarse=16, num_fine=48)\n",
+    "        self.patchwork = Patchwork(64, 2)\n",
+    "        self.head = nn.Linear(64, num_classes)\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        scatt = self.scattering(x)\n",
+    "        scatt_pool = F.adaptive_avg_pool2d(scatt, (1, 1)).squeeze(-1).squeeze(-1)\n",
+    "        emb = l2_normalize(self.jl(scatt_pool))\n",
+    "        tri = self.constellation(emb)\n",
+    "        pw = self.patchwork(tri)\n",
+    "        logits = self.head(pw)\n",
+    "        return logits, emb\n",
+    "\n",
+    "model = HierarchicalPipeline(embed_dim=128, num_classes=10).to(device)\n",
+    "print(f\"[7.3b] Model params: {sum(p.numel() for p in model.parameters()):,}\")\n",
+    "\n",
+    "result_7_3b = run_experiment(\n",
+    "    \"7.3b_Hierarchical_Constellation\",\n",
+    "    model, train_loader, val_loader, device,\n",
+    "    epochs=100, lr=1e-3, weight_decay=1e-4, patience=20\n",
+    ")\n",
+    "scoreboard.add(result_7_3b)\n",
+    ""
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 7.4a: Best Deterministic + Linear Layer\n",
+    "\n",
+    "Reference + 1 learned linear layer before constellation"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "\n",
+    "print(\"[7.4a] Best Deterministic + Linear Layer\")\n",
+    "print(\"=\"*70)\n",
+    "\n",
+    "class DeterministicLinearPipeline(nn.Module):\n",
+    "    def __init__(self, embed_dim=128, num_classes=10):\n",
+    "        super().__init__()\n",
+    "        self.scattering = Scattering2D(J=2, L=8, shape=(32, 32), max_order=2)\n",
+    "        self.jl = JLProjection(81, embed_dim)\n",
+    "        # Linear projection before constellation\n",
+    "        self.linear = nn.Linear(embed_dim, embed_dim)\n",
+    "        self.classifier = ConstellationClassifier(embed_dim, num_anchors=64, num_classes=num_classes)\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        scatt = self.scattering(x)\n",
+    "        scatt_pool = F.adaptive_avg_pool2d(scatt, (1, 1)).squeeze(-1).squeeze(-1)\n",
+    "        jl_feat = self.jl(scatt_pool)\n",
+    "        proj = self.linear(jl_feat)\n",
+    "        emb = l2_normalize(proj)\n",
+    "        logits, tri = self.classifier(emb)\n",
+    "        return logits, emb\n",
+    "\n",
+    "model = DeterministicLinearPipeline(embed_dim=128, num_classes=10).to(device)\n",
+    "print(f\"[7.4a] Model params: {sum(p.numel() for p in model.parameters()):,}\")\n",
+    "\n",
+    "result_7_4a = run_experiment(\n",
+    "    \"7.4a_DeterministicLinear\",\n",
+    "    model, train_loader, val_loader, device,\n",
+    "    epochs=100, lr=1e-3, weight_decay=1e-4, patience=20\n",
+    ")\n",
+    "scoreboard.add(result_7_4a)\n",
+    ""
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 7.4b: Deterministic + SquaredReLU + LayerNorm"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "\n",
+    "print(\"[7.4b] Deterministic + SquaredReLU + LayerNorm\")\n",
+    "print(\"=\"*70)\n",
+    "\n",
+    "class DeterministicNonlinearPipeline(nn.Module):\n",
+    "    def __init__(self, embed_dim=128, num_classes=10):\n",
+    "        super().__init__()\n",
+    "        self.scattering = Scattering2D(J=2, L=8, shape=(32, 32), max_order=2)\n",
+    "        self.jl = JLProjection(81, embed_dim)\n",
+    "        # Nonlinear projection\n",
+    "        self.proj = nn.Sequential(\n",
+    "            nn.Linear(embed_dim, embed_dim),\n",
+    "            SquaredReLU(),\n",
+    "            nn.LayerNorm(embed_dim)\n",
+    "        )\n",
+    "        self.classifier = ConstellationClassifier(embed_dim, num_anchors=64, num_classes=num_classes)\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        scatt = self.scattering(x)\n",
+    "        scatt_pool = F.adaptive_avg_pool2d(scatt, (1, 1)).squeeze(-1).squeeze(-1)\n",
+    "        jl_feat = self.jl(scatt_pool)\n",
+    "        proj = self.proj(jl_feat)\n",
+    "        emb = l2_normalize(proj)\n",
+    "        logits, tri = self.classifier(emb)\n",
+    "        return logits, emb\n",
+    "\n",
+    "model = DeterministicNonlinearPipeline(embed_dim=128, num_classes=10).to(device)\n",
+    "print(f\"[7.4b] Model params: {sum(p.numel() for p in model.parameters()):,}\")\n",
+    "\n",
+    "result_7_4b = run_experiment(\n",
+    "    \"7.4b_DeterministicNonlinear\",\n",
+    "    model, train_loader, val_loader, device,\n",
+    "    epochs=100, lr=1e-3, weight_decay=1e-4, patience=20\n",
+    ")\n",
+    "scoreboard.add(result_7_4b)\n",
+    ""
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 7.4c: Learned Projection Size Sweep\n",
+    "\n",
+    "0, 1K, 5K, 10K, 50K, 100K params"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "\n",
+    "print(\"[7.4c] Learned Projection Size Sweep\")\n",
+    "print(\"=\"*70)\n",
+    "\n",
+    "def count_proj_params(input_dim, hidden_dim, output_dim, depth=1):\n",
+    "    total = 0\n",
+    "    prev = input_dim\n",
+    "    for i in range(depth):\n",
+    "        h = hidden_dim if i < depth - 1 else output_dim\n",
+    "        total += prev * h + h\n",
+    "        prev = h\n",
+    "    return total\n",
+    "\n",
+    "class ProjectionPipeline(nn.Module):\n",
+    "    def __init__(self, embed_dim=128, num_classes=10, proj_layers=None):\n",
+    "        super().__init__()\n",
+    "        self.scattering = Scattering2D(J=2, L=8, shape=(32, 32), max_order=2)\n",
+    "        self.jl = JLProjection(81, embed_dim) if proj_layers is None else None\n",
+    "\n",
+    "        if proj_layers is None:\n",
+    "            # No projection\n",
+    "            self.proj = nn.Identity()\n",
+    "        else:\n",
+    "            # Build multi-layer projection\n",
+    "            layers = []\n",
+    "            prev_dim = 81\n",
+    "            for i, hidden_dim in enumerate(proj_layers[:-1]):\n",
+    "                layers.append(nn.Linear(prev_dim, hidden_dim))\n",
+    "                layers.append(nn.ReLU())\n",
+    "                prev_dim = hidden_dim\n",
+    "            layers.append(nn.Linear(prev_dim, embed_dim))\n",
+    "            self.proj = nn.Sequential(*layers)\n",
+    "\n",
+    "        self.classifier = ConstellationClassifier(embed_dim, num_anchors=64, num_classes=num_classes)\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        scatt = self.scattering(x)\n",
+    "        scatt_pool = F.adaptive_avg_pool2d(scatt, (1, 1)).squeeze(-1).squeeze(-1)\n",
+    "        if self.jl is not None:\n",
+    "            feat = self.jl(scatt_pool)\n",
+    "        else:\n",
+    "            feat = self.proj(scatt_pool)\n",
+    "        emb = l2_normalize(feat)\n",
+    "        logits, tri = self.classifier(emb)\n",
+    "        return logits, emb\n",
+    "\n",
+    "# 0K params: identity (no projection)\n",
+    "for target_params, proj_config in [\n",
+    "    (0, None),\n",
+    "    (1000, [128]),\n",
+    "    (5000, [256]),\n",
+    "    (10000, [512]),\n",
+    "    (50000, [512, 512]),\n",
+    "    (100000, [1024, 512]),\n",
+    "]:\n",
+    "    model = ProjectionPipeline(embed_dim=128, num_classes=10, proj_layers=proj_config).to(device)\n",
+    "    actual_params = sum(p.numel() for p in model.parameters())\n",
+    "    print(f\"[7.4c] Projection {target_params}K \u2192 actual {actual_params:,} params\")\n",
+    "\n",
+    "    result = run_experiment(\n",
+    "        f\"7.4c_ProjectionSize_{target_params}K\",\n",
+    "        model, train_loader, val_loader, device,\n",
+    "        epochs=100, lr=1e-3, weight_decay=1e-4, patience=20\n",
+    "    )\n",
+    "    scoreboard.add(result)\n",
+    ""
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Final Scoreboard"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "\n",
+    "print(\"\\n\" + \"=\"*80)\n",
+    "print(\"EXPERIMENT 7 COMPLETE: Composite Pipeline Tests\")\n",
+    "print(\"=\"*80)\n",
+    "scoreboard._print_table()\n",
+    "scoreboard.save(\"scoreboard_exp7.json\")\n",
+    ""
+   ],
+   "execution_count": null,
+   "outputs": []
+  }
+ ]
+}