#!/usr/bin/env python3 """ 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 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.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 """ import os import sys import time import argparse import numpy as np import pandas as pd import torch sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) from nqr_snn import config from nqr_snn.data.generator import ( generate_dataset_v2, generate_signal_at_snr, generate_noise_only_at_power, generate_dataset, ) from nqr_snn.data.dataset import ( NQRDatasetV2, get_balanced_loader_v2, extract_features, extract_features_batch, ) from nqr_snn.snn.model import SpikingClassifier from nqr_snn.snn.encoder import DeterministicEncoder from nqr_snn.snn.train import train_snn from nqr_snn.snn.ensemble import SNNEnsemble from nqr_snn.snn.neuron_search import run_neuron_search from nqr_snn.evaluation.metrics import full_report from nqr_snn.evaluation.plots import ( plot_roc, plot_snr_accuracy, plot_neuron_search_heatmap, 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.""" return _denoise_batch(denoiser, signals) # ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ # Helper: generate a balanced test set at one SNR level # ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ def make_snr_test_set(target_snr_db, n_per_class=200, seed=99): """Return (signals, labels) — balanced, complex128.""" rng = np.random.RandomState(seed) sigs, labs = [], [] for _ in range(n_per_class): noisy, _ = generate_signal_at_snr(target_snr_db, "white", rng) sigs.append(noisy); labs.append(1) for _ in range(n_per_class): noise = generate_noise_only_at_power("white", rng, target_power=1.0) sigs.append(noise); labs.append(0) return np.array(sigs, dtype=np.complex128), np.array(labs) # ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ # Helper: signals → features tensor (optionally denoise first) # ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ def signals_to_features(signals, denoiser=None): """(N,1024) complex128 -> (N,7,1024) float32 tensor — vectorized.""" if denoiser is not None: signals = denoise_batch(signals, denoiser) feats = extract_features_batch(signals) return torch.from_numpy(feats) # ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ # MAIN # ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ def main(): parser = argparse.ArgumentParser( description="NQR-SNN v3.2 — full pipeline" ) parser.add_argument("--quick", action="store_true", help="Quick demo mode (~90 s on CPU)") parser.add_argument("--train_size", type=int, default=None) parser.add_argument("--val_size", type=int, default=None) parser.add_argument("--ensemble_size", type=int, default=None) parser.add_argument("--max_epochs", type=int, default=None) parser.add_argument("--n_test", type=int, default=None) parser.add_argument("--skip_denoise", action="store_true") parser.add_argument("--skip_neuron_search", action="store_true") parser.add_argument("--skip_stress_test", action="store_true") parser.add_argument("--search_workers", type=int, default=None, help="Parallel workers for neuron search (default: auto)") parser.add_argument("--search_epochs", type=int, default=None) args = parser.parse_args() # -- 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) ensemble_size = args.ensemble_size or (3 if quick else config.ENSEMBLE_SIZE) max_epochs = args.max_epochs or (50 if quick else config.MAX_EPOCHS) n_test = args.n_test or (100 if quick else 300) search_epochs = args.search_epochs or (10 if quick else config.NEURON_SEARCH_EPOCHS) denoise_samples = 50 if quick else 200 stress_n = 50 if quick else 200 device = "cuda" if torch.cuda.is_available() else "cpu" print("=" * 80) print("NQR-SNN FRAMEWORK v3.2 — FULL PIPELINE") print("=" * 80) 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)}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'} | " f"stress={'ON' if not args.skip_stress_test else 'OFF'}") print("=" * 80) t0 = time.time() step_times = {} # ────────────────────────────────────────────────────────────────── # STEP 1 Denoiser Selection # ────────────────────────────────────────────────────────────────── selected_denoiser = None denoiser_name = "None" 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)") old_vs = config.VAL_SIZE config.VAL_SIZE = denoise_samples den_data = generate_dataset("low", "white", "val", seed=42) config.VAL_SIZE = old_vs selector = DenoisingSelector() denoiser_name, denoiser_r2, selected_denoiser = selector.select( den_data["noisy"][:denoise_samples], den_data["clean"][:denoise_samples], ) step_times["denoiser"] = time.time() - t1 print(f" Done: {denoiser_name} (R2={denoiser_r2:.1f}) [{step_times['denoiser']:.1f}s]") else: print("\n[1/6] DENOISER SELECTION — skipped") # ────────────────────────────────────────────────────────────────── # STEP 2 Parallel Neuron + Surrogate Search # ────────────────────────────────────────────────────────────────── neuron_model = "LIF" surrogate_fn = "ATan" if not args.skip_neuron_search: t2 = time.time() print(f"\n[2/6] PARALLEL NEURON SEARCH ({search_epochs} epochs, " f"workers={args.search_workers or 'auto'})") old_ts, old_vs = config.TRAIN_SIZE, config.VAL_SIZE s_train = min(train_size, 300) s_val = min(val_size, 100) config.TRAIN_SIZE = s_train config.VAL_SIZE = s_val config.NEURON_SEARCH_EPOCHS = search_epochs sd = generate_dataset_v2(s_train, config.TRAIN_SNR_RANGE, "white", seed=42) sd_ds = NQRDatasetV2(sd["signals"], sd["labels"], denoiser=selected_denoiser) sd_loader = get_balanced_loader_v2(sd_ds, batch_size=config.BATCH_SIZE) sv = generate_dataset_v2(s_val, config.VAL_SNR_RANGE, "white", seed=100) sv_ds = NQRDatasetV2(sv["signals"], sv["labels"], denoiser=selected_denoiser) sv_loader = get_balanced_loader_v2(sv_ds, batch_size=config.BATCH_SIZE, shuffle=False) search_result = run_neuron_search(sd_loader, sv_loader, max_workers=args.search_workers) neuron_model = search_result["model"] surrogate_fn = search_result["surrogate"] config.TRAIN_SIZE, config.VAL_SIZE = old_ts, old_vs 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} " f"(val_acc={search_result['val_accuracy']:.4f}) " f"[wall {step_times['search']:.1f}s]") else: print(f"\n[2/6] NEURON SEARCH — skipped (using {neuron_model}+{surrogate_fn})") # ────────────────────────────────────────────────────────────────── # STEP 3 Dataset Generation (with denoising) # ────────────────────────────────────────────────────────────────── t3 = time.time() print(f"\n[3/6] DATASET GENERATION (train={train_size}/class, val={val_size}/class" f"{', denoised' if selected_denoiser else ''})") config.TRAIN_SIZE = train_size config.VAL_SIZE = val_size tr = generate_dataset_v2(train_size, config.TRAIN_SNR_RANGE, "white", seed=42) tr_ds = NQRDatasetV2(tr["signals"], tr["labels"], denoiser=selected_denoiser) train_loader = get_balanced_loader_v2(tr_ds, batch_size=config.BATCH_SIZE) vl = generate_dataset_v2(val_size, config.VAL_SNR_RANGE, "white", seed=100) vl_ds = NQRDatasetV2(vl["signals"], vl["labels"], denoiser=selected_denoiser) 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 " f"SNR=[{tr['snr_dbs'].min():.0f},{tr['snr_dbs'].max():.0f}]dB " f"[{step_times['datagen']:.1f}s]") # ────────────────────────────────────────────────────────────────── # STEP 4 Ensemble Training # ────────────────────────────────────────────────────────────────── t4 = time.time() 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, ) 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]") # ────────────────────────────────────────────────────────────────── # STEP 5 SNR-Level Evaluation # ────────────────────────────────────────────────────────────────── t5 = time.time() # 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, " f"levels={config.EVAL_SNR_LEVELS})") for snr_db in config.EVAL_SNR_LEVELS: sigs, labs = make_snr_test_set(snr_db, n_per_class=n_test, seed=abs(snr_db) * 100) feat = signals_to_features(sigs, denoiser=selected_denoiser) x_seq = encoder.encode(feat).to(device) mean_p, std_p = ensemble.predict(x_seq) mean_p_np = mean_p.cpu().numpy() rep = full_report(labs, mean_p_np) rep["snr_db"] = snr_db rep["ensemble_std"] = float(std_p.cpu().mean()) snr_rows.append(rep) 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}") if snr_db in [-25, -35]: plot_roc(labs, mean_p_np, f"ROC — SNR {snr_db} dB", os.path.join(config.PLOTS_DIR, f"roc_snr{snr_db}.png")) snr_df = pd.DataFrame(snr_rows) snr_csv = os.path.join(config.RESULTS_DIR, "snn_accuracy_vs_snr.csv") snr_df.to_csv(snr_csv, index=False) plot_snr_accuracy(snr_csv) step_times["eval"] = time.time() - t5 # ────────────────────────────────────────────────────────────────── # STEP 6 Noise Injection Stress Test # ────────────────────────────────────────────────────────────────── if not args.skip_stress_test: t6 = time.time() print(f"\n[6/6] NOISE INJECTION STRESS TEST ({stress_n}/class)") st_data = generate_dataset_v2(stress_n, config.TEST_SNR_RANGE, "white", seed=99) stress_results = run_stress_test( ensemble=ensemble, test_noisy=st_data["signals"], test_labels=st_data["labels"], encoder=encoder, ) plot_noise_degradation( stress_results, os.path.join(config.PLOTS_DIR, "noise_degradation.png")) plot_ood_uncertainty( stress_results, 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]") else: print("\n[6/6] STRESS TEST — skipped") # ────────────────────────────────────────────────────────────────── # SUMMARY # ────────────────────────────────────────────────────────────────── total = time.time() - t0 print("\n" + "=" * 80) print("PIPELINE COMPLETE — RESULTS SUMMARY (v3.2)") print("=" * 80) print(f" Neuron / surrogate : {neuron_model} + {surrogate_fn}") print(f" Denoiser : {denoiser_name}") print(f" Ensemble members : {ensemble_size}") print() print(snr_df[["snr_db", "accuracy", "auc", "f1"]].to_string(index=False)) print(f"\n Mean accuracy : {snr_df['accuracy'].mean():.4f}") print(f" Min accuracy : {snr_df['accuracy'].min():.4f} " f"(SNR={snr_df.loc[snr_df['accuracy'].idxmin(), 'snr_db']:.0f} dB)") print() # Gate checks for db, tgt in [(-35, 0.95), (-40, 0.90), (-50, 0.80)]: row = snr_df[snr_df["snr_db"] == db] if len(row): a = row["accuracy"].values[0] 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(): print(f" {k:12s}: {v:6.1f}s") print(f" {'TOTAL':12s}: {total:6.1f}s ({total/60:.1f} min)") print("=" * 80) # Save all outputs os.makedirs(config.RESULTS_DIR, exist_ok=True) with open(os.path.join(config.RESULTS_DIR, "pipeline_summary.txt"), "w") as f: f.write(f"neuron_model: {neuron_model}\n") f.write(f"surrogate_fn: {surrogate_fn}\n") f.write(f"denoiser: {denoiser_name}\n") f.write(f"ensemble_size: {ensemble_size}\n") f.write(f"mean_accuracy: {snr_df['accuracy'].mean():.4f}\n") f.write(f"min_accuracy: {snr_df['accuracy'].min():.4f}\n") f.write(f"total_time_s: {total:.1f}\n") for k, v in step_times.items(): f.write(f"time_{k}_s: {v:.1f}\n") return snr_df # MANDATORY: __main__ guard for Windows spawn-based multiprocessing if __name__ == "__main__": main()