Final Evaluation for QMAML_Arnav Singhal

Quantum_MAML_for_HEP_Arnav_Singhal/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 Arnav Singhal
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

Quantum_MAML_for_HEP_Arnav_Singhal/README.md

@@ -0,0 +1,77 @@
+# QMAML: Quantum Model-Agnostic Meta-Learning
+QMAML combines classical meta-learning (MAML) with parameterized quantum circuits (PQC) to enable rapid adaptation on new tasks with few examples. This project implements and benchmarks QMAML on three datasets: MNIST, Higgs Boson, and Quark–Gluon jets. We analyze the effects of different quantum parameter initializations and demonstrate that QMAML achieves strong performance in complex scientific classification problems.
+
+Validated on:
+- **Image classification:** [MNIST](http://yann.lecun.com/exdb/mnist/) (few-shot digit recognition)  
+- **HEP data:** Higgs Boson and [Quark–Gluon jet dataset](https://arxiv.org/abs/1902.08276)  
+
+![GSoC @ ML4Sci](images/GSOC%20ML4Sci.jpg)
+
+*This project was developed as part of Google Summer of Code (GSoC) 2025 under the ML4Sci organization.*
+
+**Project Title:** Quantum Model-Agnostic Meta-Learning for Variational Quantum Algorithms for High Energy Physics Tasks at LHC  
+**Author:** Arnav Singhal  
+**Mentors:** KC Kong, Junyong Lee, Jeihee Cho  
+
+---
+
+## Task Generation (Overview)
+To ensure tasks are meaningful and balanced, we designed a **few-shot task generation pipeline**:
+- **Preprocessing:** Images normalized with train statistics.  
+- **Binning:** Dataset split by physics variables (*pt, m0*) into quantile ranges.  
+- **Sampling:** Support and query sets drawn from matched bins to avoid distribution mismatch.  
+- **Guards:** Multiple checks (distribution match, triviality checks, PCA sanity) ensure each task is non-trivial and representative.  
+
+This pipeline yields high-quality few-shot tasks for robust training and evaluation.
+
+---
+
+## Training Architecture (Overview)
+Our training setup combines **classical CNN embeddings** with **quantum circuits**:
+- **Backbone:** Small-stem ResNet-18 (with frozen batch norm) extracts stable 512-D features.  
+- **Quantum Layer:** Parameterized quantum circuit (RY + StronglyEntanglingLayers) processes embeddings into expressive quantum states.  
+- **Meta-Learning:**  
+  - **Inner Loop:** Rapid adaptation using Reptile-style updates or Q-MAML task embeddings.  
+  - **Outer Loop:** Updates both CNN and PQC components to generalize across tasks.  
+
+This hybrid design leverages classical stability with quantum adaptability.
+
+---
+
+## Dataset Links
+- [MNIST Dataset](http://yann.lecun.com/exdb/mnist/)  
+- [Quark–Gluon Jet Dataset (HEPML)](https://arxiv.org/abs/1902.08276)  
+- Higgs Boson Dataset (available on UCI/Kaggle)
+
+---
+
+## Structure
+
+The project is organized as follows:
+
+### Main Directory:
+
+• **LICENSE**: Contains the licensing information for the project.
+
+• **README**: Provides an overview of the project, its objectives, and structure.
+
+• **config.py**: Contains configuration settings and hyperparameters for training quantum meta-learning models.
+
+• **data.py**: Handles dataset loading, preprocessing, and few-shot task generation pipeline for MNIST, Higgs Boson, and Quark-Gluon jet datasets.
+
+• **train.py**: Implements the main training loops for both QMAML and Reptile-style meta-learning approaches.
+
+• **utils.py**: Contains utility functions for visualization, plotting, and performance analysis.
+
+• **models/**: A folder dedicated to model implementations. 
+
+  ○ **hybrid.py**: Implements the hybrid classical-quantum architecture combining CNN feature extractors with PQC.
+
+  ○ **pqc.py**: Contains the implementation of parameterized quantum circuits (PQC)
+
+This structure allows for the exploration of quantum meta-learning architectures across different scientific datasets (MNIST, Higgs Boson, Quark-Gluon jets). The project combines classical CNN embeddings with quantum circuits to enable rapid few-shot adaptation on new tasks.
+
+---
+
+## License
+This project is licensed under the MIT License – see the [LICENSE](LICENSE) file for details.

Quantum_MAML_for_HEP_Arnav_Singhal/final_model/config.py

@@ -0,0 +1,77 @@
+# config.py
+from dataclasses import dataclass, field
+from typing import List
+import os
+
+@dataclass
+class Config:
+    # ---- Dataset Paths (done in collab) ----
+    TRAIN_PATH: str = 'dataset path goes here'
+    TEST_PATH: str = 'dataset path goes here'
+    CHECKPOINT_DIR: str = 'dataset path goes here'
+    Q_MAML_USE: bool = True
+    LEARNER_HIDDEN: int = 256
+    FREEZE_CNN_DURING_META: bool = True
+    W0_SCALE: float = 0.01
+
+    # ---- Data / Tasking ----
+    SAMPLES: int = 90000
+    META_TASK_TYPE: List[str] = field(default_factory=lambda: ['pt', 'm0'])
+    META_BIN_COUNT: int = 6
+    SUPPORT_SIZE: int = 8
+    QUERY_SIZE: int = 8
+    MAX_META_TASKS: int = 48
+
+    # ---- Model sizes ----
+    NUM_QUBITS: int = 6
+    Q_DEPTH: int = 3
+    ENCODING_SCHEME: str = 'angle'
+    CNN_OUTPUT_DIM: int = 512
+    USE_PRETRAINED_CNN: bool = True
+
+    # ---- Task generation knobs ----
+    TG_BIN_MODE: str = "quantile"
+    TG_TARGET_JSD_MAX: float = 0.45
+    TG_JSD_MAX_TRIES: int = 24
+    TG_NUM_TASKS_PER_BIN: int = 6
+    TG_TRAIN_SEED: int = 42
+    TG_TEST_SEED: int = 43
+
+    # ---- Caps & Guards ----
+    TG_CAP_PT: float = 0.65
+    TG_CAP_M0: float = 0.60
+    GUARD_LOGREG_AUC_MIN: float = 0.58
+
+    # ---- PCA Guard ----
+    PCA_COMPONENTS: int = 128
+    PCA_GUARD_CAL_TRIALS: int = 60
+    PCA_GUARD_PERCENTILE: float = 0.60
+    PCA_GUARD_CLAMP_MIN: float = 0.56
+    PCA_GUARD_CLAMP_MAX: float = 0.90
+
+    # ---- Warm-up / Adaptive gates ----
+    WARMUP_ACCEPT_N: int = 24
+    WARMUP_PCA_DELTA: float = 0.03
+    ADAPTIVE_M0_LOOSE_CAP: float = 0.50
+    ADAPTIVE_M0_MIN_ACCEPT: int = 3
+
+    # ---- Training ----
+    USE_ANALYTIC_GRADIENTS: bool = True
+    INNER_STEPS: int = 12
+    INNER_LR: float = 0.02
+    OUTER_LR: float = 5e-3
+    EPOCHS: int = 15
+    BATCH_SIZE: int = 24
+    EVAL_METRICS: bool = True
+    SAVE_BEST_MODEL: bool = True
+
+    # ---- Pretraining ----
+    PRETRAIN_QMAML: bool = True
+    USE_QMAML_PRETRAINED_INIT: bool = True
+    QMAML_INIT_TAG: str = "qmaml_init"
+
+
+config = Config()
+os.makedirs(config.CHECKPOINT_DIR, exist_ok=True)
+
+__all__ = ["Config", "config"]

Quantum_MAML_for_HEP_Arnav_Singhal/final_model/data.py

@@ -0,0 +1,57 @@
+import h5py
+import numpy as np
+import torch
+from torch.utils.data import Dataset
+from typing import Tuple
+from config import Config, config
+
+class JetDataset(Dataset):
+    def __init__(self, X: np.ndarray, y: np.ndarray, pt: np.ndarray, m0: np.ndarray):
+        self.X  = X.astype(np.float32, copy=False)    # (N,H,W,C)
+        self.y  = y.astype(np.int64,   copy=False)    # (N,)
+        self.pt = pt.astype(np.float32, copy=False)
+        self.m0 = m0.astype(np.float32, copy=False)
+
+    def __len__(self) -> int:
+        return len(self.y)
+
+    def __getitem__(self, idx: int):
+        return self.X[idx], int(self.y[idx])
+
+def _describe_dataset(name: str, ds: JetDataset):
+    N = len(ds)
+    yc = np.bincount(ds.y, minlength=2)
+    print(f"[{name}] N={N}  y-counts={yc.tolist()}  frac(class1)={yc[1]/max(1,yc.sum()):.3f}")
+    for feat in ("pt","m0"):
+        a = getattr(ds, feat)
+        print(f"  {feat}: min={a.min():.3f} p10={np.percentile(a,10):.3f} "
+              f"p50={np.percentile(a,50):.3f} p90={np.percentile(a,90):.3f} max={a.max():.3f}")
+    Xs = ds.X[:5]
+    print(f"  X sample: shape={Xs.shape}, dtype={Xs.dtype}, min={Xs.min():.3f}, max={Xs.max():.3f}")
+
+def load_datasets(cfg: Config) -> Tuple[JetDataset, JetDataset]:
+    def _load(path: str, take: int) -> JetDataset:
+        with h5py.File(path, 'r') as f:
+            X  = f['X_jets'][:take]
+            y  = f['y'][:take]
+            pt = f['pt'][:take]
+            m0 = f['m0'][:take]
+        return JetDataset(X, y, pt, m0)
+
+    train = _load(cfg.TRAIN_PATH, cfg.SAMPLES)
+    test  = _load(cfg.TEST_PATH, cfg.SAMPLES)
+    _describe_dataset("TRAIN(raw)", train)
+    _describe_dataset("TEST(raw)",  test)
+    return train, test
+
+# Train-only per-channel normalization (applied to both splits)
+train_dataset, test_dataset = load_datasets(config)
+GLOBAL_NORM = {
+    "mean": train_dataset.X.mean(axis=(0,1,2), keepdims=True),
+    "std":  train_dataset.X.std(axis=(0,1,2),  keepdims=True) + 1e-6
+}
+
+
+__all__ = [
+    "JetDataset", "load_datasets", "train_dataset", "test_dataset", "GLOBAL_NORM"
+]

Quantum_MAML_for_HEP_Arnav_Singhal/final_model/main.py

@@ -0,0 +1,143 @@
+import torch
+import torch.nn as nn
+import numpy as np
+from typing import Dict, List
+
+from config import config
+from data import train_dataset, test_dataset
+from tasks import meta_tasks, test_meta_tasks
+
+from models import CNNFeatureExtractor, PQCModel, HybridModel, freeze_bn
+from meta import inner_loop_adaptation, outer_loop_meta_update, outer_loop_qmaml
+from utils import plot_training_results, plot_comparison
+
+# Helper: flatten PQC-only params
+
+def _pqc_only_flat(pqc: nn.Module) -> torch.Tensor:
+    """Flatten PQC weights + fc params (stable shape across inner loop)."""
+    return torch.cat([
+        pqc.weights.detach().flatten().cpu(),
+        pqc.fc.weight.detach().flatten().cpu(),
+        pqc.fc.bias.detach().flatten().cpu(),
+    ])
+
+# Warmup PQC head
+
+def warmup_pqc_head(model: nn.Module, meta_tasks: List[Dict], steps: int = 100, lr: float = 1e-2):
+    """Light warm-up of PQC head before full training."""
+    opt = torch.optim.SGD(
+        [{"params": [model.pqc.weights], "lr": lr},
+         {"params": model.pqc.fc.parameters(), "lr": lr}],
+        momentum=0.9
+    )
+    loss_fn = nn.CrossEntropyLoss()
+    model.train()
+    for i in range(steps):
+        task = meta_tasks[i % len(meta_tasks)]
+        opt.zero_grad()
+        logits = model(task["support_X"])
+        loss = loss_fn(logits, task["support_y"])
+        loss.backward()
+        opt.step()
+
+# Main Experiment Loop
+
+def run_all_inits():
+    initialization_types = ["gaussian", "uniform", "qmaml_learner", "zero", "pi"]
+    results: Dict[str, Dict[str, List[float]]] = {}
+
+    for init_type in initialization_types:
+        print(f"\n=== Testing initialization: {init_type} ===")
+        torch.manual_seed(42)
+        np.random.seed(42)
+
+        # --- Build CNN ---
+        cnn_extractor = CNNFeatureExtractor(config.CNN_OUTPUT_DIM, config.NUM_QUBITS)
+        freeze_bn(cnn_extractor)
+
+        # --- Build PQC ---
+        pqc_init = "zero" if init_type == "qmaml_learner" else init_type
+        pqc_model = PQCModel(
+            config.NUM_QUBITS,
+            config.Q_DEPTH,
+            init_type=pqc_init,
+            bound_angles=True,
+            verbose=True
+        )
+        hybrid_model = HybridModel(cnn_extractor, pqc_model)
+
+        # --- Probes ---
+        _sx = meta_tasks[0]["support_X"][:8]
+        _sy = meta_tasks[0]["support_y"][:8]
+        _qx = meta_tasks[0]["query_X"]
+
+        # 7A) Gradient flow
+        for p in hybrid_model.parameters():
+            if p.grad is not None:
+                p.grad.zero_()
+        _loss = nn.CrossEntropyLoss()(hybrid_model(_sx), _sy)
+        _loss.backward()
+        def _gnorm(p): return None if (p.grad is None) else float(p.grad.norm().item())
+        try:
+            _cnn_fc0 = cnn_extractor.backbone.fc[0].weight
+            print("[Probe 7A]", init_type,
+                  " grad||pqc.weights:", _gnorm(pqc_model.weights),
+                  " grad||pqc.fc.weight:", _gnorm(pqc_model.fc.weight),
+                  " grad||cnn.backbone.fc[0].weight:", _gnorm(_cnn_fc0))
+        except Exception:
+            print("[Probe 7A]", init_type,
+                  " grad||pqc.weights:", _gnorm(pqc_model.weights),
+                  " grad||pqc.fc.weight:", _gnorm(pqc_model.fc.weight))
+
+        # 7B) Inner loop PQC-only test
+        w_before = _pqc_only_flat(pqc_model)
+        _ = inner_loop_adaptation(
+            hybrid_model,
+            meta_tasks[0]["support_X"], meta_tasks[0]["support_y"],
+            inner_steps=2, inner_lr=0.01
+        )
+        w_after = _pqc_only_flat(pqc_model)
+        print("[Probe 7B]", init_type, "||Δθ_PQC|| =", torch.norm(w_after - w_before).item())
+
+        # 7C) CNN feature stats
+        with torch.no_grad():
+            _f = cnn_extractor(meta_tasks[0]["support_X"])
+        print("[Probe 7C]", init_type, " feat mean:", _f.mean().item(), " feat std:", _f.std().item())
+
+        # 7D) Prediction distribution
+        with torch.no_grad():
+            _p = torch.softmax(hybrid_model(_qx[:16]), dim=1)[:, 1]
+        print("[Probe 7D]", init_type, " p1[min,mean,max]:",
+              _p.min().item(), _p.mean().item(), _p.max().item())
+
+        # --- Warm-up PQC head (all inits) ---
+        warmup_pqc_head(hybrid_model, meta_tasks, steps=100, lr=1e-2)
+
+        # --- Train ---
+        if init_type == "qmaml_learner":
+            training_results = outer_loop_qmaml(
+                hybrid_model,
+                meta_tasks,
+                test_meta_tasks,
+                config.OUTER_LR,
+                config.EVAL_METRICS,
+                ckpt_name="best_qmaml_learner.pth",
+            )
+        else:
+            training_results = outer_loop_meta_update(
+                hybrid_model,
+                meta_tasks,
+                test_meta_tasks,
+                config.OUTER_LR,
+                config.EVAL_METRICS,
+                ckpt_name=f"best_{init_type}.pth",
+            )
+
+        results[init_type] = training_results
+        plot_training_results(training_results, init_type)
+
+    plot_comparison(results)
+
+
+if __name__ == "__main__":
+    run_all_inits()