""" Heavy Task Feasibility Study: Parameter Dimension & Data Hardness Scaling. Protocol Version: HEAVY-TASK-PSO-V6 1.0.0 Feasibility probe to test retained PSO methods (G0, G5, G6, G8) along two axes: 1. Larger Parameter Dimension (WideCNN ~55k params vs CompactCNN 9,098 params) 2. Harder Data (FashionMNIST vs MNIST) Workload Matrix (4 workloads x 4 methods = 16 screen cells): - mnist_compact (MNIST + CompactCNN 9,098 params, control) - mnist_wide (MNIST + WideCNN ~55k params, parameter scaling axis) - fashion_compact (FashionMNIST + CompactCNN 9,098 params, data hardness axis) - fashion_wide (FashionMNIST + WideCNN ~55k params, combined axis) Methods: G0, G5, G6, G8 """ from __future__ import annotations import argparse import copy from dataclasses import dataclass, asdict import hashlib import json import csv import math import sys import time from pathlib import Path from typing import Any, Dict, List, Optional, Tuple, Union import numpy as np import torch import torch.nn as nn from sklearn.model_selection import train_test_split import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt from matplotlib.lines import Line2D # Ensure test directory and repo root are in Python path REPO_ROOT = Path(__file__).resolve().parent.parent if Path(__file__).resolve().parent.name != "PSO" else Path(__file__).resolve().parent TEST_DIR = REPO_ROOT / "test" if str(TEST_DIR) not in sys.path: sys.path.insert(0, str(TEST_DIR)) if str(REPO_ROOT) not in sys.path: sys.path.insert(0, str(REPO_ROOT)) from benchmark_suite import ( calc_stats, compute_model_fingerprint, get_hardware_provenance, resolve_execution_device, save_json_atomic, sync_device, ) from deep_pso_methods import ( CompactCNN, make_compact_cnn, build_nested_stratified_subsets, evaluate_probabilistic_metrics, get_model_probabilities, ) from deep_pso_v6 import ( V6GeometryConfig, get_v6_geometry_table, V6LatentTransform, run_v6_pso, run_g8_optimizer, ) from pso import __version__ as pso_version PROTOCOL_VERSION = "HEAVY-TASK-PSO-V6 1.0.0" FEASIBILITY_NLL_REDUCTION = 0.20 FEASIBILITY_ACCURACY_GAIN_PP = 20.0 # ===================================================================== # 1. Architecture Definitions & Deterministic Factories # ===================================================================== class WideCNN(nn.Module): """ WideCNN architecture (~55,338 parameters): Conv1 (1->16, 3x3, pad=1), ReLU, MaxPool2d(2,2) Conv2 (16->32, 3x3, pad=1), ReLU, MaxPool2d(2,2) Flatten -> 32x7x7 = 1568 Linear (1568 -> 32), ReLU Linear (32 -> 10) """ def __init__(self): super().__init__() self.conv1 = nn.Conv2d(1, 16, kernel_size=3, padding=1) self.relu1 = nn.ReLU() self.pool1 = nn.MaxPool2d(2, 2) self.conv2 = nn.Conv2d(16, 32, kernel_size=3, padding=1) self.relu2 = nn.ReLU() self.pool2 = nn.MaxPool2d(2, 2) self.flatten = nn.Flatten() self.fc1 = nn.Linear(1568, 32) self.relu3 = nn.ReLU() self.fc2 = nn.Linear(32, 10) def forward(self, x: torch.Tensor) -> torch.Tensor: if x.dim() == 2 and x.shape[1] == 784: x = x.view(-1, 1, 28, 28) out = self.pool1(self.relu1(self.conv1(x))) out = self.pool2(self.relu2(self.conv2(out))) out = self.flatten(out) out = self.relu3(self.fc1(out)) return self.fc2(out) def make_wide_cnn(seed: int = 41) -> nn.Module: """Deterministic factory for WideCNN by seed.""" torch.manual_seed(seed) return WideCNN() def create_model(model_name: str, seed: int = 41) -> nn.Module: """Factory function creating a model instance by name and seed.""" name = model_name.lower() if name in ("compact_cnn", "compactcnn"): return make_compact_cnn(seed) elif name in ("wide_cnn", "widecnn"): return make_wide_cnn(seed) else: raise ValueError(f"Unknown model architecture name: '{model_name}'. Expected 'compact_cnn' or 'wide_cnn'.") # ===================================================================== # 2. Generic Train-Only Data Preparation (MNIST & FashionMNIST) # ===================================================================== def prepare_heavy_task_data( dataset_name: str, split_seed: int = 20260902, cache_dir: Optional[Path] = None, ) -> Tuple[ torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, Dict[int, torch.Tensor], str, Dict[str, Any] ]: """ Train-only data preparation for MNIST or FashionMNIST using exclusively train=True. Never constructs train=False. Preserves exact split seed (20260902), search-only normalization, and nested 2k/10k/50k index stratification. """ if cache_dir is None: cache_dir = Path("result/cache") cache_dir.mkdir(parents=True, exist_ok=True) ds_lower = dataset_name.lower() if ds_lower in ("mnist", "mnist_compact", "mnist_wide"): from torchvision.datasets import MNIST raw_train = MNIST(root=str(cache_dir), train=True, download=True) canonical_name = "MNIST" elif ds_lower in ("fashion_mnist", "fashion", "fashion_compact", "fashion_wide"): from torchvision.datasets import FashionMNIST raw_train = FashionMNIST(root=str(cache_dir), train=True, download=True) canonical_name = "FashionMNIST" else: raise ValueError(f"Unsupported dataset name: '{dataset_name}'. Must be 'mnist' or 'fashion_mnist'.") x_train_raw = raw_train.data.float() / 255.0 # (60000, 28, 28) y_train_raw = raw_train.targets.long() # Stratified split: 50,000 search set and 10,000 validation set indices = np.arange(len(y_train_raw)) search_idx, val_idx = train_test_split( indices, train_size=50000, test_size=10000, stratify=y_train_raw.numpy(), random_state=split_seed, ) x_search_raw = x_train_raw[search_idx] y_search = y_train_raw[search_idx] x_val_raw = x_train_raw[val_idx] y_val = y_train_raw[val_idx] # Fit mean and std on 50k search subset ONLY mean_val = float(x_search_raw.mean()) std_val = float(x_search_raw.std()) x_search_norm = ((x_search_raw - mean_val) / std_val).unsqueeze(1) # (50000, 1, 28, 28) x_val_norm = ((x_val_raw - mean_val) / std_val).unsqueeze(1) # (10000, 1, 28, 28) # Nested stratified subsets inside 50k search set: 2k inside 10k inside 50k nested_subsets = build_nested_stratified_subsets( y_search=y_search, subset_sizes=[2000, 10000, 50000], subset_seed=split_seed, ) # Data fingerprint over search and validation splits (no test split) h = hashlib.sha256() for t in (x_search_norm, x_val_norm, y_search, y_val): h.update(t.detach().cpu().numpy().tobytes()) data_fp = h.hexdigest()[:16] split_h = hashlib.sha256() split_h.update(search_idx.tobytes()) split_h.update(val_idx.tobytes()) split_fp = split_h.hexdigest()[:16] provenance = { "dataset_name": canonical_name, "input_shape": [1, 28, 28], "normalization_scope": "search_train_50000_only", "train_mean": round(mean_val, 6), "train_std": round(std_val, 6), "search_samples": 50000, "val_samples": 10000, "test_samples": 0, "official_test_data_loaded": False, "official_test_evaluations": 0, "split_seed": split_seed, "split_fingerprint": split_fp, "data_fingerprint": data_fp, } return ( x_search_norm, y_search, x_val_norm, y_val, nested_subsets, data_fp, provenance ) # ===================================================================== # 3. Workload Configurations & Method Specifications # ===================================================================== @dataclass(frozen=True) class WorkloadConfig: workload_id: str dataset_name: str model_name: str description: str WORKLOADS: Dict[str, WorkloadConfig] = { "mnist_compact": WorkloadConfig( workload_id="mnist_compact", dataset_name="mnist", model_name="compact_cnn", description="MNIST dataset with CompactCNN (9,098 params control)", ), "mnist_wide": WorkloadConfig( workload_id="mnist_wide", dataset_name="mnist", model_name="wide_cnn", description="MNIST dataset with WideCNN (~55k params larger model axis)", ), "fashion_compact": WorkloadConfig( workload_id="fashion_compact", dataset_name="fashion_mnist", model_name="compact_cnn", description="FashionMNIST dataset with CompactCNN (harder data axis)", ), "fashion_wide": WorkloadConfig( workload_id="fashion_wide", dataset_name="fashion_mnist", model_name="wide_cnn", description="FashionMNIST dataset with WideCNN (harder data + larger model axis)", ), } HEAVY_METHODS = ["G0", "G5", "G6", "G8"] def get_heavy_geometry_table() -> Dict[str, V6GeometryConfig]: """Returns the subset of V6 geometry configurations used for heavy task study (G0, G5, G6, G8).""" full_table = get_v6_geometry_table() return {m: full_table[m] for m in HEAVY_METHODS} # ===================================================================== # 4. Untrained Baseline Evaluation # ===================================================================== def evaluate_untrained_baseline( base_model: nn.Module, x_sub: torch.Tensor, y_sub: torch.Tensor, x_val: torch.Tensor, y_val: torch.Tensor, device: torch.device, ) -> Dict[str, float]: """ Evaluates an untrained base model on validation set and objective subset to establish baseline performance for feasibility classification. """ model = copy.deepcopy(base_model).to(device) val_probs = get_model_probabilities(model, x_val.to(device), device) val_metrics = evaluate_probabilistic_metrics(val_probs, y_val.to(device)) # Evaluate objective subset (e.g. 2k or 10k) model.eval() loss_fn = nn.CrossEntropyLoss(reduction="sum") x_sub_dev = x_sub.to(device) y_sub_dev = y_sub.to(device) num_sub = len(y_sub_dev) with torch.inference_mode(): total_loss = 0.0 correct = 0 for b_start in range(0, num_sub, 1000): xb = x_sub_dev[b_start:b_start + 1000] yb = y_sub_dev[b_start:b_start + 1000] logits = model(xb) total_loss += float(loss_fn(logits, yb).item()) correct += int((logits.argmax(dim=1) == yb).sum().item()) obj_loss = total_loss / num_sub obj_acc = (correct / num_sub) * 100.0 return { "val_nll": val_metrics["nll"], "val_accuracy": val_metrics["accuracy"], "val_brier": val_metrics["brier"], "val_ece": val_metrics["ece"], "val_margin": val_metrics["margin"], "objective_loss": round(obj_loss, 6), "objective_accuracy": round(obj_acc, 4), } # ===================================================================== # 5. Method Selection & Feasibility Classification Logic # ===================================================================== def select_best_normalized_method( screen_results_for_workload: List[Dict[str, Any]] ) -> str: """ Selects the best normalized custom method among G0, G5, G6 for a workload based on screen validation NLL ascending, with validation accuracy descending tiebreak. """ candidates = [] for result in screen_results_for_workload: if result["method_id"] not in ("G0", "G5", "G6"): continue val_loss = result.get("val_selected_loss") val_acc = result.get("val_selected_acc") if ( val_loss is not None and val_acc is not None and math.isfinite(val_loss) and math.isfinite(val_acc) ): candidates.append(result) if not candidates: raise ValueError("No finite normalized method (G0, G5, G6) result found for selection.") return min( candidates, key=lambda result: (result["val_selected_loss"], -result["val_selected_acc"]), )["method_id"] def evaluate_feasibility( confirmed_runs: List[Dict[str, Any]], baseline_val_nll: float, baseline_val_acc: float, ) -> Dict[str, Any]: """ Classifies feasibility per workload and selected method: - execution_feasible: True iff every run completed with finite metrics. - optimization_feasible: True iff mean validation NLL is at least 20% below baseline and mean validation accuracy is at least 20 percentage points above baseline. """ is_execution_feasible = True nll_vals = [] acc_vals = [] for run in confirmed_runs: val_nll = run.get("val_selected_loss") val_acc = run.get("val_selected_acc") if val_nll is None or val_acc is None: is_execution_feasible = False break if not (math.isfinite(val_nll) and math.isfinite(val_acc)): is_execution_feasible = False break nll_vals.append(val_nll) acc_vals.append(val_acc) target_nll_thresh = baseline_val_nll * (1.0 - FEASIBILITY_NLL_REDUCTION) target_acc_thresh = baseline_val_acc + FEASIBILITY_ACCURACY_GAIN_PP if not is_execution_feasible or len(nll_vals) == 0: return { "execution_feasible": False, "optimization_feasible": False, "baseline_val_nll": baseline_val_nll, "target_val_nll_threshold": round(target_nll_thresh, 6), "baseline_val_acc": baseline_val_acc, "target_val_acc_threshold": round(target_acc_thresh, 4), "mean_val_nll": None, "mean_val_acc": None, } mean_nll = float(np.mean(nll_vals)) mean_acc = float(np.mean(acc_vals)) is_optimization_feasible = (mean_nll <= target_nll_thresh) and (mean_acc >= target_acc_thresh) return { "execution_feasible": True, "optimization_feasible": is_optimization_feasible, "baseline_val_nll": baseline_val_nll, "target_val_nll_threshold": round(target_nll_thresh, 6), "baseline_val_acc": baseline_val_acc, "target_val_acc_threshold": round(target_acc_thresh, 4), "mean_val_nll": round(mean_nll, 6), "mean_val_acc": round(mean_acc, 4), } # ===================================================================== # 6. Screening Stage Runner (16 Workload-Method Cells) # ===================================================================== def run_heavy_task_screen( workloads: Dict[str, WorkloadConfig], methods: List[str], particles: int = 12, epochs: int = 40, seed: int = 91, device: torch.device = torch.device("cpu"), cache_dir: Optional[Path] = None, ) -> Tuple[List[Dict[str, Any]], Dict[str, Dict[str, Any]], Dict[str, Any]]: """ Screens all 16 workload-method cells at fixed 2k subset, seed 91, 12p x 40e. """ screen_results = [] untrained_baselines = {} data_cache = {} geom_table = get_heavy_geometry_table() for wl_id, wl_cfg in workloads.items(): if wl_cfg.dataset_name not in data_cache: x_search, y_search, x_val, y_val, nested_subsets, data_fp, provenance = prepare_heavy_task_data( dataset_name=wl_cfg.dataset_name, cache_dir=cache_dir, ) data_cache[wl_cfg.dataset_name] = { "x_search": x_search, "y_search": y_search, "x_val": x_val, "y_val": y_val, "nested_subsets": nested_subsets, "data_fp": data_fp, "provenance": provenance, } else: d = data_cache[wl_cfg.dataset_name] x_search, y_search = d["x_search"], d["y_search"] x_val, y_val = d["x_val"], d["y_val"] nested_subsets = d["nested_subsets"] base_model = create_model(wl_cfg.model_name, seed=41) param_count = sum(p.numel() for p in base_model.parameters()) x_2k = x_search[nested_subsets[2000]] y_2k = y_search[nested_subsets[2000]] if wl_id not in untrained_baselines: untrained_baselines[wl_id] = evaluate_untrained_baseline( base_model, x_2k, y_2k, x_val, y_val, device ) for m_id in methods: geom_config = geom_table[m_id] if m_id in ("G0", "G5", "G6"): transform = V6LatentTransform(base_model, geom_config, device) res = run_v6_pso( transform=transform, base_model=base_model, x_search=x_search, y_search=y_search, x_val=x_val, y_val=y_val, nested_subsets=nested_subsets, schedule_str=f"2000:{epochs}", epochs=epochs, swarm_size=particles, seed=seed, device=device, geom_config=geom_config, val_check_interval=10, ) # Analytical core swarm-state bytes for custom methods (Z, V, P, M, V_sq) core_swarm_state_bytes = 5 * particles * param_count * 4 elif m_id == "G8": res = run_g8_optimizer( base_model=base_model, x_2k=x_2k, y_2k=y_2k, x_val=x_val, y_val=y_val, epochs=epochs, swarm_size=particles, seed=seed, device=device, ) # Analytical core swarm-state bytes for public Optimizer (5*particles + 1 global best) core_swarm_state_bytes = (5 * particles + 1) * param_count * 4 else: raise ValueError(f"Unknown method ID: {m_id}") opt_time = max(res["optimization_wall_time_sec"], 1e-6) total_samples = res["total_sample_evaluations"] throughput_sps = round(total_samples / opt_time, 2) val_nll = res["val_selected_loss"] val_acc = res["val_selected_acc"] g_loss = res["gbest_loss"] g_acc = res["gbest_acc"] finite_metrics = (val_nll, val_acc, g_loss, g_acc) is_finite = all( val is not None and math.isfinite(val) for val in finite_metrics ) cell_record = { "workload_id": wl_id, "method_id": m_id, "dataset_name": wl_cfg.dataset_name, "model_name": wl_cfg.model_name, "parameter_count": param_count, "subset_size": 2000, "particles": particles, "epochs": epochs, "seed": seed, "gbest_loss": g_loss, "gbest_acc": g_acc, "gbest_val_loss": res.get("gbest_val_loss"), "gbest_val_acc": res.get("gbest_val_acc"), "val_selected_particle_idx": res.get("val_selected_particle_idx"), "val_selected_loss": val_nll, "val_selected_acc": val_acc, "val_metrics": res.get("val_metrics"), "wall_time_sec": res["wall_time_sec"], "optimization_wall_time_sec": res["optimization_wall_time_sec"], "validation_wall_time_sec": res["validation_wall_time_sec"], "total_queries": res["total_queries"], "total_sample_evaluations": res["total_sample_evaluations"], "validation_evaluations": res["validation_evaluations"], "official_test_evaluations": 0, "core_swarm_state_bytes": core_swarm_state_bytes, "throughput_samples_per_sec": throughput_sps, "is_finite": is_finite, } screen_results.append(cell_record) # Workload summary metadata workload_metadata = {} for wl_id, wl_cfg in workloads.items(): bm = create_model(wl_cfg.model_name, seed=41) param_count = sum(p.numel() for p in bm.parameters()) model_fp = compute_model_fingerprint(bm) d = data_cache[wl_cfg.dataset_name] workload_metadata[wl_id] = { "workload_id": wl_id, "dataset_name": wl_cfg.dataset_name, "model_name": wl_cfg.model_name, "parameter_count": param_count, "model_fingerprint": model_fp, "data_fingerprint": d["data_fp"], "split_fingerprint": d["provenance"]["split_fingerprint"], "description": wl_cfg.description, } return screen_results, untrained_baselines, workload_metadata # ===================================================================== # 7. Confirmation Stage Runner (Seeds 101-103, Fixed 10k, 12p x 80e) # ===================================================================== def run_heavy_task_confirm( workloads: Dict[str, WorkloadConfig], selected_methods: Dict[str, List[str]], # workload_id -> [G8, best_normalized] particles: int = 12, epochs: int = 80, seeds: List[int] = (101, 102, 103), split_seed: int = 20260902, device: torch.device = torch.device("cpu"), cache_dir: Optional[Path] = None, ) -> Dict[str, Dict[str, Any]]: """ Confirms selected methods (G8 + best normalized) for each workload at fixed 10k, 12p x 80e, seeds 101-103. Aggregates mean and sample SD across seeds. """ confirmation_results = {} geom_table = get_heavy_geometry_table() data_cache = {} for wl_id, wl_cfg in workloads.items(): if wl_cfg.dataset_name not in data_cache: x_search, y_search, x_val, y_val, nested_subsets, data_fp, provenance = prepare_heavy_task_data( dataset_name=wl_cfg.dataset_name, split_seed=split_seed, cache_dir=cache_dir, ) data_cache[wl_cfg.dataset_name] = { "x_search": x_search, "y_search": y_search, "x_val": x_val, "y_val": y_val, "nested_subsets": nested_subsets, "data_fp": data_fp, "provenance": provenance, } else: d = data_cache[wl_cfg.dataset_name] x_search, y_search = d["x_search"], d["y_search"] x_val, y_val = d["x_val"], d["y_val"] nested_subsets = d["nested_subsets"] data_fp = d["data_fp"] provenance = d["provenance"] base_model = create_model(wl_cfg.model_name, seed=41) param_count = sum(p.numel() for p in base_model.parameters()) x_10k = x_search[nested_subsets[10000]] y_10k = y_search[nested_subsets[10000]] wl_confirmations = {} methods_to_confirm = selected_methods[wl_id] for m_id in methods_to_confirm: geom_config = geom_table[m_id] per_seed_runs = [] for s in seeds: if m_id in ("G0", "G5", "G6"): transform = V6LatentTransform(base_model, geom_config, device) res = run_v6_pso( transform=transform, base_model=base_model, x_search=x_search, y_search=y_search, x_val=x_val, y_val=y_val, nested_subsets=nested_subsets, schedule_str=f"10000:{epochs}", epochs=epochs, swarm_size=particles, seed=s, device=device, geom_config=geom_config, val_check_interval=10, ) core_bytes = 5 * particles * param_count * 4 elif m_id == "G8": res = run_g8_optimizer( base_model=base_model, x_2k=x_10k, # Passes 10k search tensor for 10k fit y_2k=y_10k, x_val=x_val, y_val=y_val, epochs=epochs, swarm_size=particles, seed=s, device=device, ) core_bytes = (5 * particles + 1) * param_count * 4 else: raise ValueError(f"Unknown method ID: {m_id}") opt_time = max(res["optimization_wall_time_sec"], 1e-6) sps = round(res["total_sample_evaluations"] / opt_time, 2) seed_record = { "seed": s, "val_selected_loss": res["val_selected_loss"], "val_selected_acc": res["val_selected_acc"], "val_metrics": res.get("val_metrics"), "gbest_loss": res["gbest_loss"], "gbest_acc": res["gbest_acc"], "wall_time_sec": res["wall_time_sec"], "optimization_wall_time_sec": res["optimization_wall_time_sec"], "validation_wall_time_sec": res["validation_wall_time_sec"], "total_queries": res["total_queries"], "total_sample_evaluations": res["total_sample_evaluations"], "validation_evaluations": res["validation_evaluations"], "official_test_evaluations": 0, "core_swarm_state_bytes": core_bytes, "throughput_samples_per_sec": sps, } per_seed_runs.append(seed_record) # Compute aggregated mean & sample SD stats acc_list = [r["val_selected_acc"] for r in per_seed_runs] nll_list = [r["val_selected_loss"] for r in per_seed_runs] brier_list = [r["val_metrics"]["brier"] for r in per_seed_runs if r.get("val_metrics")] ece_list = [r["val_metrics"]["ece"] for r in per_seed_runs if r.get("val_metrics")] g_loss_list = [r["gbest_loss"] for r in per_seed_runs] g_acc_list = [r["gbest_acc"] for r in per_seed_runs] wall_list = [r["wall_time_sec"] for r in per_seed_runs] opt_wall_list = [r["optimization_wall_time_sec"] for r in per_seed_runs] sps_list = [r["throughput_samples_per_sec"] for r in per_seed_runs] stats = { "val_acc": calc_stats(acc_list), "val_nll": calc_stats(nll_list), "val_brier": calc_stats(brier_list) if brier_list else None, "val_ece": calc_stats(ece_list) if ece_list else None, "gbest_loss": calc_stats(g_loss_list), "gbest_acc": calc_stats(g_acc_list), "wall_time_sec": calc_stats(wall_list), "optimization_wall_time_sec": calc_stats(opt_wall_list), "throughput_samples_per_sec": calc_stats(sps_list), } wl_confirmations[m_id] = { "workload_id": wl_id, "method_id": m_id, "subset_size": 10000, "particles": particles, "epochs": epochs, "seeds": list(seeds), "split_seed": split_seed, "data_fingerprint": data_fp, "split_fingerprint": provenance["split_fingerprint"], "provenance": provenance, "stats": stats, "per_seed_runs": per_seed_runs, } confirmation_results[wl_id] = wl_confirmations return confirmation_results # ===================================================================== # 8. CSV & Plot Artifact Generators # ===================================================================== def save_csv_summary(payload: Dict[str, Any], csv_path: Path): """Saves a clean, concise CSV summary of screen, confirmation, and feasibility results.""" csv_path.parent.mkdir(parents=True, exist_ok=True) with open(csv_path, "w", newline="", encoding="utf-8") as f: writer = csv.writer(f) writer.writerow(["section", "workload", "method", "metric", "value"]) writer.writerow(["protocol", "global", "all", "version", payload["protocol_version"]]) writer.writerow(["protocol", "global", "all", "official_test_data_loaded", payload["official_test_data_loaded"]]) writer.writerow(["protocol", "global", "all", "official_test_evaluations", payload["official_test_evaluations"]]) # Untrained Baselines baselines = payload.get("untrained_baselines", {}) for wl_id, base in baselines.items(): writer.writerow(["baseline", wl_id, "untrained", "val_nll", base["val_nll"]]) writer.writerow(["baseline", wl_id, "untrained", "val_accuracy", base["val_accuracy"]]) # Screen Results screen_res = payload.get("screen_results", []) for r in screen_res: wl = r["workload_id"] m = r["method_id"] writer.writerow(["screen", wl, m, "val_nll", r["val_selected_loss"]]) writer.writerow(["screen", wl, m, "val_acc", r["val_selected_acc"]]) writer.writerow(["screen", wl, m, "gbest_loss", r["gbest_loss"]]) writer.writerow(["screen", wl, m, "gbest_acc", r["gbest_acc"]]) writer.writerow(["screen", wl, m, "throughput_sps", r["throughput_samples_per_sec"]]) # Confirmation Results confirm_res = payload.get("confirmation_results", {}) for wl_id, m_dict in confirm_res.items(): for m_id, conf in m_dict.items(): st = conf["stats"] writer.writerow(["confirm", wl_id, m_id, "val_acc_mean", st["val_acc"]["mean"]]) writer.writerow(["confirm", wl_id, m_id, "val_acc_std", st["val_acc"]["std"]]) writer.writerow(["confirm", wl_id, m_id, "val_nll_mean", st["val_nll"]["mean"]]) writer.writerow(["confirm", wl_id, m_id, "val_nll_std", st["val_nll"]["std"]]) writer.writerow(["confirm", wl_id, m_id, "wall_time_sec_mean", st["wall_time_sec"]["mean"]]) # Feasibility Evaluations feas_evals = payload.get("feasibility_evaluations", {}) for wl_id, m_feas in feas_evals.items(): for m_id, fe in m_feas.items(): writer.writerow(["feasibility", wl_id, m_id, "execution_feasible", fe["execution_feasible"]]) writer.writerow(["feasibility", wl_id, m_id, "optimization_feasible", fe["optimization_feasible"]]) writer.writerow(["feasibility", wl_id, m_id, "mean_val_nll", fe["mean_val_nll"]]) writer.writerow(["feasibility", wl_id, m_id, "mean_val_acc", fe["mean_val_acc"]]) def generate_feasibility_plot(payload: Dict[str, Any], plot_path: Path): """ Generates a clear two-panel visualization: Panel 1: Validation NLL screen comparison across all 16 workload-method cells. Panel 2: Confirmation Accuracy vs Parameter Count for each workload, distinguishing dataset (MNIST vs FashionMNIST), model size, and method. """ plot_path.parent.mkdir(parents=True, exist_ok=True) fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15.5, 6)) # Panel 1: Screen Validation NLL screen_results = payload.get("screen_results", []) workload_order = ["mnist_compact", "mnist_wide", "fashion_compact", "fashion_wide"] method_order = ["G0", "G5", "G6", "G8"] cell_dict = {(r["workload_id"], r["method_id"]): r["val_selected_loss"] for r in screen_results} x_indices = np.arange(len(workload_order)) width = 0.18 colors = {"G0": "#1f77b4", "G5": "#ff7f0e", "G6": "#2ca02c", "G8": "#d62728"} for idx, m_id in enumerate(method_order): vals = [cell_dict.get((wl_id, m_id), 0.0) for wl_id in workload_order] ax1.bar(x_indices + (idx - 1.5) * width, vals, width, label=m_id, color=colors[m_id]) ax1.set_xticks(x_indices) ax1.set_xticklabels(["MNIST\nCompact", "MNIST\nWide", "Fashion\nCompact", "Fashion\nWide"], fontsize=9) ax1.set_ylabel("Validation NLL (Screen fixed2k, lower is better)") ax1.set_title("Screening Stage: Validation NLL (16 Cells)") ax1.legend(title="Method") ax1.grid(True, linestyle="--", alpha=0.5) # Panel 2: Confirmation Accuracy vs Parameter Count confirm_results = payload.get("confirmation_results", {}) workload_info = payload.get("workloads", {}) markers = {"mnist": "o", "fashion_mnist": "s"} dataset_x_factors = {"mnist": 0.94, "fashion_mnist": 1.06} method_x_factors = {"G8": 0.985, "G0": 1.015, "G5": 1.015, "G6": 1.015} parameter_counts = sorted({ metadata["parameter_count"] for metadata in workload_info.values() }) for parameter_count in parameter_counts: ax2.axvline(parameter_count, color="#bbbbbb", linewidth=0.8, linestyle=":") for wl_id, m_dict in confirm_results.items(): wl_meta = workload_info.get(wl_id, {}) ds_name = wl_meta.get("dataset_name", "mnist").lower() param_count = wl_meta.get("parameter_count", 9098) marker = markers.get(ds_name, "o") for m_id, conf in m_dict.items(): acc_mean = conf["stats"]["val_acc"]["mean"] acc_std = conf["stats"]["val_acc"]["std"] color = colors.get(m_id, "#333333") display_x = ( param_count * dataset_x_factors.get(ds_name, 1.0) * method_x_factors.get(m_id, 1.0) ) ax2.errorbar( [display_x], [acc_mean], yerr=[acc_std], fmt=marker, color=color, linestyle="none", capsize=5, markersize=8, label="_nolegend_", ) ax2.set_xlabel("Model Parameter Count D (log scale; points horizontally offset)") ax2.set_ylabel("Validation Accuracy % (Confirmation fixed10k, mean ± SD)") ax2.set_title("Confirmation: G8 + Screen-Selected Normalized Method") ax2.set_xticks(parameter_counts) ax2.set_xticklabels([f"{count:,}" for count in parameter_counts]) ax2.grid(True, axis="y", linestyle="--", alpha=0.5) ax2.set_xscale("log") dataset_handles = [ Line2D( [0], [0], marker=marker, color="#333333", linestyle="none", markersize=8, label=label, ) for label, marker in (("MNIST", "o"), ("FashionMNIST", "s")) ] confirmed_method_ids = sorted({ method_id for methods in confirm_results.values() for method_id in methods }) method_handles = [ Line2D( [0], [0], marker="o", color=colors[method_id], linestyle="none", markersize=8, label=method_id, ) for method_id in confirmed_method_ids ] dataset_legend = ax2.legend( handles=dataset_handles, title="Dataset marker", fontsize=8, loc="upper left", bbox_to_anchor=(1.01, 1.0), ) ax2.add_artist(dataset_legend) ax2.legend( handles=method_handles, title="Method color", fontsize=8, loc="upper left", bbox_to_anchor=(1.01, 0.68), ) plt.tight_layout() plt.savefig(plot_path, dpi=150) plt.close(fig) # ===================================================================== # 9. Main Orchestration Runner & CLI Entrypoint # ===================================================================== def run_heavy_task_study(args: argparse.Namespace) -> Dict[str, Any]: """Orchestrates screen, confirmation, feasibility evaluation, and artifact generation.""" start_time_all = time.time() device = resolve_execution_device(args.device) cache_dir = Path(args.cache_dir) out_dir = Path(args.out_dir) plot_dir = Path(args.plot_dir) out_dir.mkdir(parents=True, exist_ok=True) plot_dir.mkdir(parents=True, exist_ok=True) screen_particles = args.screen_particles screen_epochs = args.screen_epochs screen_seed = args.screen_seeds[0] if args.screen_seeds else 91 confirm_particles = args.confirm_particles confirm_epochs = args.confirm_epochs confirm_seeds = args.confirm_seeds screen_design = { "methods": list(HEAVY_METHODS), "subset_size": 2000, "particles": screen_particles, "epochs": screen_epochs, "seed": screen_seed, "cells": len(WORKLOADS) * len(HEAVY_METHODS), } method_configs = { method_id: asdict(config) for method_id, config in get_heavy_geometry_table().items() } screen_results = [] untrained_baselines = {} workload_metadata = {} normalized_selections = {} if args.stage in ("screen", "all"): print(f"[{PROTOCOL_VERSION}] Starting Screening Stage (16 Cells at fixed2k, {screen_particles}p x {screen_epochs}e, seed {screen_seed})...") screen_results, untrained_baselines, workload_metadata = run_heavy_task_screen( workloads=WORKLOADS, methods=HEAVY_METHODS, particles=screen_particles, epochs=screen_epochs, seed=screen_seed, device=device, cache_dir=cache_dir, ) # Select best normalized method for each workload for wl_id in WORKLOADS: wl_screen = [r for r in screen_results if r["workload_id"] == wl_id] best_norm = select_best_normalized_method(wl_screen) normalized_selections[wl_id] = best_norm print(f" Workload '{wl_id}': G8 + selected best normalized method '{best_norm}'") if args.stage == "screen": # Save screen-only artifact screen_payload = { "protocol_version": PROTOCOL_VERSION, "timestamp": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()), "hardware": get_hardware_provenance(device), "pso_version": pso_version, "official_test_data_loaded": False, "official_test_evaluations": 0, "study_design": { "screen": screen_design, "selection": { "candidates": ["G0", "G5", "G6"], "ranking": ["validation_nll_ascending", "validation_accuracy_descending"], "confirmation_reference": "G8", }, "validation": { "source": "official_training_split_only", "split": "50000_search_10000_validation_stratified", "split_seed": 20260902, }, }, "method_configs": method_configs, "workloads": workload_metadata, "untrained_baselines": untrained_baselines, "screen_results": screen_results, "normalized_method_selections": normalized_selections, } save_json_atomic(screen_payload, out_dir / "pso_v6_heavy_tasks_screen.json") print(f"Screening complete. Saved to {out_dir / 'pso_v6_heavy_tasks_screen.json'}") return screen_payload elif args.stage == "confirm": screen_artifact_path = args.screen_artifact or (out_dir / "pso_v6_heavy_tasks_screen.json") if not screen_artifact_path.exists(): screen_artifact_path = out_dir / "pso_v6_heavy_tasks.json" if not screen_artifact_path.exists(): raise FileNotFoundError(f"Screen artifact not found at {screen_artifact_path}. Run --stage screen or --stage all first.") with open(screen_artifact_path, "r", encoding="utf-8") as f: screen_payload = json.load(f) screen_results = screen_payload["screen_results"] untrained_baselines = screen_payload["untrained_baselines"] workload_metadata = screen_payload["workloads"] normalized_selections = screen_payload["normalized_method_selections"] screen_design = screen_payload["study_design"]["screen"] # Confirmation Stage selected_methods_to_confirm = {} for wl_id in WORKLOADS: best_norm = normalized_selections[wl_id] selected_methods_to_confirm[wl_id] = ["G8", best_norm] print(f"[{PROTOCOL_VERSION}] Starting Confirmation Stage (Seeds {confirm_seeds}, fixed10k, {confirm_particles}p x {confirm_epochs}e)...") confirm_results = run_heavy_task_confirm( workloads=WORKLOADS, selected_methods=selected_methods_to_confirm, particles=confirm_particles, epochs=confirm_epochs, seeds=confirm_seeds, device=device, cache_dir=cache_dir, ) # Feasibility Evaluations feasibility_evals = {} for wl_id, m_dict in confirm_results.items(): base = untrained_baselines[wl_id] wl_feas = {} for m_id, conf in m_dict.items(): feas = evaluate_feasibility( confirmed_runs=conf["per_seed_runs"], baseline_val_nll=base["val_nll"], baseline_val_acc=base["val_accuracy"], ) wl_feas[m_id] = feas feasibility_evals[wl_id] = wl_feas # Aggregate Resource Totals tot_queries = 0 tot_samples = 0 tot_opt_wall = 0.0 tot_val_wall = 0.0 tot_run_wall = 0.0 for r in screen_results: tot_queries += r["total_queries"] tot_samples += r["total_sample_evaluations"] tot_opt_wall += r["optimization_wall_time_sec"] tot_val_wall += r["validation_wall_time_sec"] tot_run_wall += r["wall_time_sec"] for wl_id, m_dict in confirm_results.items(): for m_id, conf in m_dict.items(): for run in conf["per_seed_runs"]: tot_queries += run["total_queries"] tot_samples += run["total_sample_evaluations"] tot_opt_wall += run["optimization_wall_time_sec"] tot_val_wall += run["validation_wall_time_sec"] tot_run_wall += run["wall_time_sec"] resource_totals = { "total_queries": tot_queries, "total_sample_evaluations": tot_samples, "summed_optimization_wall_time_sec": round(tot_opt_wall, 4), "summed_validation_wall_time_sec": round(tot_val_wall, 4), "summed_recorded_run_wall_time_sec": round(tot_run_wall, 4), "elapsed_current_process_wall_time_sec": round(time.time() - start_time_all, 4), } final_payload = { "protocol_version": PROTOCOL_VERSION, "timestamp": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()), "hardware": get_hardware_provenance(device), "pso_version": pso_version, "official_test_data_loaded": False, "official_test_evaluations": 0, "study_design": { "screen": screen_design, "confirmation": { "methods_per_workload": selected_methods_to_confirm, "subset_size": 10000, "particles": confirm_particles, "epochs": confirm_epochs, "seeds": list(confirm_seeds), }, "selection": { "candidates": ["G0", "G5", "G6"], "ranking": ["validation_nll_ascending", "validation_accuracy_descending"], "confirmation_reference": "G8", }, "feasibility_contract": { "execution": "all_confirmed_runs_complete_with_finite_metrics", "minimum_validation_nll_reduction_fraction": FEASIBILITY_NLL_REDUCTION, "minimum_validation_accuracy_gain_percentage_points": FEASIBILITY_ACCURACY_GAIN_PP, "requires_both_optimization_thresholds": True, }, "validation": { "source": "official_training_split_only", "split": "50000_search_10000_validation_stratified", "split_seed": 20260902, "official_test_split_used": False, }, }, "method_configs": method_configs, "workloads": workload_metadata, "untrained_baselines": untrained_baselines, "screen_results": screen_results, "normalized_method_selections": normalized_selections, "confirmation_results": confirm_results, "feasibility_evaluations": feasibility_evals, "resource_totals": resource_totals, } # Persist JSON, CSV, and plot json_path = out_dir / "pso_v6_heavy_tasks.json" csv_path = out_dir / "pso_v6_heavy_tasks.csv" plot_path = plot_dir / "pso_v6_heavy_tasks.png" save_json_atomic(final_payload, json_path) save_csv_summary(final_payload, csv_path) generate_feasibility_plot(final_payload, plot_path) print(f"[{PROTOCOL_VERSION}] Study complete!") print(f" JSON: {json_path}") print(f" CSV: {csv_path}") print(f" Plot: {plot_path}") return final_payload def build_parser() -> argparse.ArgumentParser: parser = argparse.ArgumentParser(description="MNIST / FashionMNIST PSO Heavy Task Feasibility Study") parser.add_argument("--stage", choices=["screen", "confirm", "all"], default="all", help="Study stage to run") parser.add_argument("--device", type=str, default=None, help="Device (cpu, mps, cuda)") parser.add_argument("--screen-particles", type=int, default=12, help="Screening particle count") parser.add_argument("--screen-epochs", type=int, default=40, help="Screening epoch count") parser.add_argument("--confirm-particles", type=int, default=12, help="Confirmation particle count") parser.add_argument("--confirm-epochs", type=int, default=80, help="Confirmation epoch count") parser.add_argument("--screen-seeds", type=int, nargs="+", default=[91], help="Screening seed") parser.add_argument("--confirm-seeds", type=int, nargs="+", default=[101, 102, 103], help="Confirmation seeds") parser.add_argument("--cache-dir", type=str, default="result/cache", help="Data cache directory") parser.add_argument("--out-dir", type=str, default="benchmark_results", help="Output directory") parser.add_argument("--plot-dir", type=str, default="history_plt", help="Plot directory") parser.add_argument("--screen-artifact", type=Path, default=None, help="Screen JSON artifact for confirm-only stage") return parser def main(): parser = build_parser() args = parser.parse_args() run_heavy_task_study(args) if __name__ == "__main__": main()