v3.2: wire neuron search results into ENSEMBLE_CONFIGS, use trained encoder at eval, consolidate denoise_batch

run_full_pipeline.py

@@ -1,24 +1,25 @@
 #!/usr/bin/env python3
 """
-NQR-SNN Framework v3.1 — Full Pipeline
+NQR-SNN Framework v3.2 — Full Pipeline
 =======================================
 Complete end-to-end pipeline with all 6 stages:
 
   1. Denoiser Selection      — LM vs SSA vs Wavelet (gate on low-SNR white)
-  2. Parallel Neuron Search  — 5 neuron × 3 surrogate = 15 configs via ProcessPoolExecutor
+  2. Parallel Neuron Search  — 5 neuron x 3 surrogate = 15 configs via ProcessPoolExecutor
   3. Dataset Generation      — SNR-controlled signals, denoised, 7-channel feature extraction
   4. Ensemble Training       — N-member heterogeneous CNN+SNN ensemble with TET loss
   5. SNR-Level Evaluation    — Accuracy / AUC / F1 at each dB level with denoising
   6. Noise Injection Stress  — Gaussian / S&P / RFI / weight / OOD frequency shift
 
-v3.1 fixes: dead neuron bug (threshold+LayerNorm), differentiable encoder, dropout.
+v3.2 fixes:
+- Neuron search results now wire into ENSEMBLE_CONFIGS (was dead code)
+- Uses trained encoder at evaluation (matches training)
+- Consolidated denoise_batch helper (no more inline loops)
 
 Usage:
-    python run_full_pipeline.py --quick                      # fast demo (~90s CPU)
-    python run_full_pipeline.py                              # full run
-    python run_full_pipeline.py --skip_denoise               # skip denoiser stage
-    python run_full_pipeline.py --skip_neuron_search          # skip neuron search, use LIF+ATan
-    python run_full_pipeline.py --search_workers 8            # force 8 parallel workers
+    python run_full_pipeline.py --quick          # fast demo (~90s CPU)
+    python run_full_pipeline.py                  # full run
+    python run_full_pipeline.py --skip_denoise   # skip denoiser stage
 """
 
 import os
@@ -50,23 +51,17 @@ from nqr_snn.evaluation.plots import (
     plot_noise_degradation, plot_ood_uncertainty,
 )
 from nqr_snn.denoising.selector import DenoisingSelector
+from nqr_snn.denoising import denoise_batch as _denoise_batch
 from nqr_snn.noise_injection.stress_test import run_stress_test
 
 
 # ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
 # Helper: apply denoiser to a batch of complex signals
+# v3.2: Now delegates to nqr_snn.denoising.denoise_batch (was duplicated)
 # ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
 def denoise_batch(signals: np.ndarray, denoiser) -> np.ndarray:
     """Apply denoiser to every signal in a (N, 1024) complex128 array."""
-    out = np.zeros_like(signals)
-    for i in range(len(signals)):
-        if hasattr(denoiser, 'fit_one'):      # LMDenoiser
-            out[i] = denoiser.fit_one(signals[i])
-        elif hasattr(denoiser, 'denoise'):     # SSA / Wavelet
-            out[i] = denoiser.denoise(signals[i])
-        else:
-            out[i] = signals[i]
-    return out
+    return _denoise_batch(denoiser, signals)
 
 
 # ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
@@ -89,10 +84,9 @@ def make_snr_test_set(target_snr_db, n_per_class=200, seed=99):
 # Helper: signals → features tensor (optionally denoise first)
 # ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
 def signals_to_features(signals, denoiser=None):
-    """(N,1024) complex128 → (N,7,1024) float32 tensor — v3.0 vectorized."""
+    """(N,1024) complex128 -> (N,7,1024) float32 tensor — vectorized."""
     if denoiser is not None:
         signals = denoise_batch(signals, denoiser)
-    # v3.0: Use vectorized batch extraction
     feats = extract_features_batch(signals)
     return torch.from_numpy(feats)
 
@@ -102,7 +96,7 @@ def signals_to_features(signals, denoiser=None):
 # ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
 def main():
     parser = argparse.ArgumentParser(
-        description="NQR-SNN v2.0 — full pipeline: denoise → neuron search → train → evaluate → stress test"
+        description="NQR-SNN v3.2 — full pipeline"
     )
     parser.add_argument("--quick", action="store_true",
                         help="Quick demo mode (~90 s on CPU)")
@@ -119,7 +113,7 @@ def main():
     parser.add_argument("--search_epochs", type=int, default=None)
     args = parser.parse_args()
 
-    # ── Resolve sizes depending on mode ──
+    # -- Resolve sizes depending on mode --
     quick = args.quick
     train_size      = args.train_size      or (500  if quick else 3000)
     val_size        = args.val_size        or (150  if quick else 1000)
@@ -133,15 +127,14 @@ def main():
     device = "cuda" if torch.cuda.is_available() else "cpu"
 
     print("=" * 80)
-    print("NQR-SNN FRAMEWORK v3.1 — FULL PIPELINE (SNN BUG FIXED)")
+    print("NQR-SNN FRAMEWORK v3.2 — FULL PIPELINE")
     print("=" * 80)
-    print(f"  v3.1 fixes: dead neuron (threshold+LayerNorm), differentiable encoder, dropout")
     print(f"  Device           : {device}")
     print(f"  Mode             : {'QUICK' if quick else 'FULL'}")
     print(f"  Train / Val      : {train_size} / {val_size}  per class")
     print(f"  Ensemble members : {ensemble_size}")
     print(f"  Max epochs       : {max_epochs}")
-    print(f"  Search epochs    : {search_epochs}  ({len(config.NEURON_MODELS)}×{len(config.SURROGATE_FUNCTIONS)} configs)")
+    print(f"  Search epochs    : {search_epochs}  ({len(config.NEURON_MODELS)}x{len(config.SURROGATE_FUNCTIONS)} configs)")
     print(f"  Test per SNR     : {n_test}")
     print(f"  Steps            : denoise={'ON' if not args.skip_denoise else 'OFF'} | "
           f"search={'ON' if not args.skip_neuron_search else 'OFF'} | "
@@ -159,7 +152,7 @@ def main():
 
     if not args.skip_denoise:
         t1 = time.time()
-        print(f"\n[1/6] DENOISER SELECTION (LM → SSA → Wavelet on {denoise_samples} low-SNR white samples)")
+        print(f"\n[1/6] DENOISER SELECTION (LM -> SSA -> Wavelet on {denoise_samples} low-SNR white samples)")
 
         old_vs = config.VAL_SIZE
         config.VAL_SIZE = denoise_samples
@@ -172,7 +165,7 @@ def main():
             den_data["clean"][:denoise_samples],
         )
         step_times["denoiser"] = time.time() - t1
-        print(f"  ✓ Selected: {denoiser_name}  (R²={denoiser_r2:.1f})  [{step_times['denoiser']:.1f}s]")
+        print(f"  Done: {denoiser_name}  (R2={denoiser_r2:.1f})  [{step_times['denoiser']:.1f}s]")
     else:
         print("\n[1/6] DENOISER SELECTION — skipped")
 
@@ -210,12 +203,40 @@ def main():
 
         step_times["search"] = time.time() - t2
 
+        # v3.2 FIX: Wire neuron search results into ENSEMBLE_CONFIGS
+        # Previously the search result was stored but ENSEMBLE_CONFIGS stayed hardcoded,
+        # making the entire search stage dead code. Now the top-K diverse configs from
+        # the search populate the ensemble, ensuring the search actually affects training.
+        search_df = search_result.get("all_results")
+        if search_df is not None and len(search_df) > 0:
+            valid = search_df[search_df["val_accuracy"] > 0].sort_values(
+                "val_accuracy", ascending=False
+            )
+            if len(valid) >= 3:
+                new_configs = []
+                for _, row in valid.iterrows():
+                    new_configs.append((
+                        row["neuron_model"],
+                        row["surrogate_fn"],
+                        config.SNN_TAU,  # default tau; vary below
+                    ))
+                # Add tau diversity: cycle through [1.5, 2.0, 2.5, 3.0] for top configs
+                tau_cycle = [1.5, 2.0, 2.5, 3.0]
+                ensemble_configs = []
+                for i, (nm, sf, _) in enumerate(new_configs):
+                    tau = tau_cycle[i % len(tau_cycle)]
+                    ensemble_configs.append((nm, sf, tau))
+                    if len(ensemble_configs) >= config.ENSEMBLE_SIZE:
+                        break
+                config.ENSEMBLE_CONFIGS = ensemble_configs
+                print(f"  Updated ENSEMBLE_CONFIGS from search: {len(ensemble_configs)} configs")
+
         # Heatmap
         csv = os.path.join(config.RESULTS_DIR, "neuron_search.csv")
         if os.path.exists(csv):
             plot_neuron_search_heatmap(csv)
 
-        print(f"  ✓ Winner: {neuron_model} + {surrogate_fn} "
+        print(f"  Winner: {neuron_model} + {surrogate_fn} "
               f"(val_acc={search_result['val_accuracy']:.4f})  "
               f"[wall {step_times['search']:.1f}s]")
     else:
@@ -240,7 +261,7 @@ def main():
     val_loader = get_balanced_loader_v2(vl_ds, batch_size=config.BATCH_SIZE, shuffle=False)
 
     step_times["datagen"] = time.time() - t3
-    print(f"  ✓ {len(tr_ds)} train, {len(vl_ds)} val  "
+    print(f"  {len(tr_ds)} train, {len(vl_ds)} val  "
           f"SNR=[{tr['snr_dbs'].min():.0f},{tr['snr_dbs'].max():.0f}]dB  "
           f"[{step_times['datagen']:.1f}s]")
 
@@ -248,25 +269,26 @@ def main():
     # STEP 4  Ensemble Training
     # ──────────────────────────────────────────────────────────────────
     t4 = time.time()
-    print(f"\n[4/6] ENSEMBLE TRAINING  ({ensemble_size} members × {max_epochs} epochs, "
+    print(f"\n[4/6] ENSEMBLE TRAINING  ({ensemble_size} members x {max_epochs} epochs, "
           f"neuron={neuron_model}+{surrogate_fn})")
 
     ensemble = SNNEnsemble(
         neuron_model=neuron_model, surrogate_fn=surrogate_fn,
         ensemble_size=ensemble_size, device=device,
-        heterogeneous=True,  # v3.0: diverse neuron/surrogate/tau
+        heterogeneous=True,
     )
     ensemble.train_all(train_loader, val_loader, max_epochs=max_epochs)
     ensemble.load_checkpoints()
 
     step_times["train"] = time.time() - t4
-    print(f"  ✓ Training done  [{step_times['train']:.1f}s]")
+    print(f"  Training done  [{step_times['train']:.1f}s]")
 
     # ──────────────────────────────────────────────────────────────────
     # STEP 5  SNR-Level Evaluation
     # ──────────────────────────────────────────────────────────────────
     t5 = time.time()
-    encoder = DeterministicEncoder()
+    # v3.2: Use ensemble's encoder (LearnableTemporalEncoder with trained weights)
+    encoder = ensemble.encoder
     snr_rows = []
 
     print(f"\n[5/6] SNR EVALUATION  ({n_test}/class per level, "
@@ -285,7 +307,7 @@ def main():
         rep["ensemble_std"] = float(std_p.cpu().mean())
         snr_rows.append(rep)
 
-        print(f"  SNR={snr_db:4d} dB │ acc={rep['accuracy']:.4f}  "
+        print(f"  SNR={snr_db:4d} dB | acc={rep['accuracy']:.4f}  "
               f"auc={rep['auc']:.4f}  f1={rep['f1']:.4f}  "
               f"tpr={rep['tpr']:.4f}  tnr={rep['tnr']:.4f}")
 
@@ -322,7 +344,7 @@ def main():
             os.path.join(config.PLOTS_DIR, "ood_uncertainty.png"))
 
         step_times["stress"] = time.time() - t6
-        print(f"  ✓ Stress test done  [{step_times['stress']:.1f}s]")
+        print(f"  Stress test done  [{step_times['stress']:.1f}s]")
     else:
         print("\n[6/6] STRESS TEST — skipped")
 
@@ -331,7 +353,7 @@ def main():
     # ──────────────────────────────────────────────────────────────────
     total = time.time() - t0
     print("\n" + "=" * 80)
-    print("PIPELINE COMPLETE — RESULTS SUMMARY (v3.1 FIXED)")
+    print("PIPELINE COMPLETE — RESULTS SUMMARY (v3.2)")
     print("=" * 80)
     print(f"  Neuron / surrogate : {neuron_model} + {surrogate_fn}")
     print(f"  Denoiser           : {denoiser_name}")
@@ -348,7 +370,8 @@ def main():
         row = snr_df[snr_df["snr_db"] == db]
         if len(row):
             a = row["accuracy"].values[0]
-            print(f"  {'✅' if a >= tgt else '⚠️ '} SNR={db} dB : {a:.4f}  (target ≥{tgt})")
+            status = "PASS" if a >= tgt else "FAIL"
+            print(f"  [{status}] SNR={db} dB : {a:.4f}  (target >={tgt})")
 
     print(f"\n  Step timings:")
     for k, v in step_times.items():