#!/usr/bin/env python3 """ PSO Stage Plugin Method Comparison Tool. Compares PSO movement methods (original, inertia, constriction, fips, clpso, bare_bones, adaptive_moment, local_best, quantum) across benchmark datasets (xor, iris, mnist) with reproducible model initialization, dataset splits, and clean console output & JSON result reporting. """ import argparse import json import sys import time from typing import Any import torch import torch.nn as nn from cli import add_pso_args, parse_method_options from pso import Optimizer from pso.plugins import available_plugins def print_available_methods(): """Prints available movement method plugins and metadata provenance.""" movement_plugins = available_plugins("movement") print("Available PSO Movement Method Plugins:") print("=" * 80) for name, meta in movement_plugins.items(): print(f" Stage Key : {name}") print(f" Title : {meta.title}") print(f" Source (DOI) : {meta.source or 'N/A'}") print(f" Fidelity : {meta.fidelity}") print(f" Needs Grad : {meta.gradient_required}") print("-" * 80) def get_xor_workload(seed: int): """Builds identical XOR dataset and PyTorch model state for a given seed.""" torch.manual_seed(seed) x_train = torch.tensor( [[0.0, 0.0], [0.0, 1.0], [1.0, 0.0], [1.0, 1.0]], dtype=torch.float32 ) y_train = torch.tensor([[0.0], [1.0], [1.0], [0.0]], dtype=torch.float32) model = nn.Sequential( nn.Linear(2, 4), nn.Tanh(), nn.Linear(4, 1), ) loss_fn = nn.BCEWithLogitsLoss() return x_train, y_train, None, model, loss_fn, "binary" def get_iris_workload(seed: int): """Builds identical Iris dataset and PyTorch model state for a given seed.""" from sklearn.datasets import load_iris from sklearn.model_selection import train_test_split from sklearn.preprocessing import StandardScaler torch.manual_seed(seed) iris = load_iris() X = iris.data.astype("float32") y = iris.target.astype("int64") x_train, x_test, y_train, y_test = train_test_split( X, y, test_size=0.2, shuffle=True, stratify=y, random_state=seed ) scaler = StandardScaler() x_train = scaler.fit_transform(x_train) x_test = scaler.transform(x_test) x_tr = torch.tensor(x_train, dtype=torch.float32) y_tr = torch.tensor(y_train, dtype=torch.int64) x_te = torch.tensor(x_test, dtype=torch.float32) y_te = torch.tensor(y_test, dtype=torch.int64) model = nn.Sequential( nn.Linear(4, 10), nn.ReLU(), nn.Linear(10, 10), nn.ReLU(), nn.Linear(10, 3), ) loss_fn = nn.CrossEntropyLoss() return x_tr, y_tr, (x_te, y_te), model, loss_fn, "multiclass" def get_mnist_workload(seed: int): """Builds identical PCA32 MNIST dataset and PyTorch model state for a given seed.""" from sklearn.decomposition import PCA from torchvision.datasets import MNIST torch.manual_seed(seed) train_ds = MNIST(root="./data", train=True, download=True) test_ds = MNIST(root="./data", train=False, download=True) x_tr_raw = (train_ds.data[:3000].float() / 255.0).reshape(3000, -1).numpy() y_tr = train_ds.targets[:3000].long() x_te_raw = (test_ds.data[:1000].float() / 255.0).reshape(1000, -1).numpy() y_te = test_ds.targets[:1000].long() pca = PCA(n_components=32, whiten=True, random_state=seed) x_tr_pca = pca.fit_transform(x_tr_raw) x_te_pca = pca.transform(x_te_raw) x_tr = torch.tensor(x_tr_pca, dtype=torch.float32) x_te = torch.tensor(x_te_pca, dtype=torch.float32) model = nn.Linear(32, 10) loss_fn = nn.CrossEntropyLoss() return x_tr, y_tr, (x_te, y_te), model, loss_fn, "multiclass" DATASET_LOADERS = { "xor": get_xor_workload, "iris": get_iris_workload, "mnist": get_mnist_workload, } def main(): parser = argparse.ArgumentParser( description="PSO Stage Plugin Method Comparison Surface" ) parser.add_argument( "--list-methods", action="store_true", help="List available movement methods and exit", ) parser.add_argument( "--dataset", choices=["xor", "iris", "mnist"], default="xor", help="Target dataset workload (default: %(default)s)", ) parser.add_argument( "--methods", nargs="+", default=["all"], help="Movement method keys to evaluate or 'all' (default: all)", ) parser.add_argument( "--seeds", type=int, nargs="+", default=[42], help="Random seed list (default: 42)", ) parser.add_argument( "--json-path", "--output-json", "--json", dest="json_path", type=str, default=None, help="Optional JSON file path to save detailed evaluation metrics", ) # Add standard stage selector and hyperparameter options add_pso_args(parser, defaults={"n_particles": 20, "epochs": 30}) args = parser.parse_args() if args.list_methods: print_available_methods() sys.exit(0) # Determine movement methods to test available_m_plugins = available_plugins("movement") if "all" in args.methods or "ALL" in args.methods: methods_to_test = list(available_m_plugins.keys()) else: methods_to_test = [] for m in args.methods: if m not in available_m_plugins: raise ValueError( f"Unknown movement method '{m}'. Available: {list(available_m_plugins.keys())}" ) methods_to_test.append(m) loader = DATASET_LOADERS[args.dataset] eval_stage = args.evaluation fitness_size = args.fitness_size if eval_stage == "fixed_subset" else None if eval_stage == "fixed_subset" and fitness_size is None: if args.dataset == "xor": fitness_size = 4 elif args.dataset == "iris": fitness_size = 100 elif args.dataset == "mnist": fitness_size = 2000 refine_stage = args.refinement refinement_epochs = args.refinement_epochs if refine_stage == "adam" else 0 results: list[dict[str, Any]] = [] parsed_method_opts = parse_method_options(args.method_options) has_val = args.dataset != "xor" print("=" * 80) print(f"PSO Movement Method Comparison on '{args.dataset}' Dataset") print( f"Stages: initialization='{args.initialization}', evaluation='{eval_stage}', " f"convergence='{args.convergence}', refinement='{refine_stage}'" ) print( f"Parameters: particles={args.n_particles}, epochs={args.epochs}, seeds={args.seeds}" ) print("=" * 80) if has_val: print( f"{'Method':<18} {'Seed':<6} {'Val Loss':<12} {'Val Accuracy':<12} {'Val MSE':<12} {'Time (s)':<10}" ) else: print( f"{'Method':<18} {'Seed':<6} {'Loss':<12} {'Accuracy':<12} {'MSE':<12} {'Time (s)':<10}" ) print("-" * 80) for method_key in methods_to_test: for seed in args.seeds: x_tr, y_tr, val_data, model, loss_fn, task = loader(seed) eff_fitness_size = ( min(fitness_size, x_tr.shape[0]) if fitness_size is not None else None ) c0 = args.c0 c1 = args.c1 w_min = args.w_min w_max = args.w_max neg_swarm = ( args.negative_swarm if method_key not in ("fips", "clpso", "bare_bones") else 0.0 ) mut_swarm = args.mutation_swarm if method_key != "bare_bones" else 0.0 vel_ratio = ( args.velocity_limit_ratio if method_key != "bare_bones" else None ) kwargs: dict[str, Any] = { "model": model, "loss": loss_fn, "task": task, "method": method_key, "initialization": args.initialization, "evaluation": eval_stage, "convergence": args.convergence, "refinement": refine_stage, "method_options": parsed_method_opts, "n_particles": args.n_particles, "c0": c0, "c1": c1, "w_min": w_min, "w_max": w_max, "negative_swarm": neg_swarm, "mutation_swarm": mut_swarm, "particle_min": args.particle_min, "particle_max": args.particle_max, "velocity_limit_ratio": vel_ratio, "boundary_strategy": args.boundary_strategy, "initial_position_noise": args.initial_position_noise, "seed": seed, "device": args.device, "fitness_size": eff_fitness_size, "convergence_patience": args.convergence_patience, "convergence_min_delta": args.convergence_min_delta, "convergence_monitor": args.convergence_monitor, "refinement_epochs": refinement_epochs, "refinement_lr": args.refinement_lr, } if method_key == "adaptive_moment": if args.moment_blend is not None: kwargs["moment_blend"] = args.moment_blend elif "moment_blend" in parsed_method_opts: kwargs["moment_blend"] = float(parsed_method_opts["moment_blend"]) else: kwargs["moment_blend"] = 0.25 if args.moment_beta1 is not None: kwargs["moment_beta1"] = args.moment_beta1 if args.moment_beta2 is not None: kwargs["moment_beta2"] = args.moment_beta2 if args.moment_step_size is not None: kwargs["moment_step_size"] = args.moment_step_size if args.moment_epsilon is not None: kwargs["moment_epsilon"] = args.moment_epsilon opt = Optimizer(**kwargs) start_t = time.perf_counter() best_score = opt.fit( x_tr, y_tr, epochs=args.epochs, batch_size=args.batch_size, fitness_size=eff_fitness_size, renewal=args.renewal, validation_data=val_data, output_dir=None, refinement_epochs=refinement_epochs, refinement_lr=args.refinement_lr, ) elapsed_t = time.perf_counter() - start_t tr_loss, tr_acc, tr_mse = best_score if val_data is not None: val_x, val_y = val_data val_score = opt.evaluate(val_x, val_y) eval_loss, eval_acc, eval_mse = val_score score_src = "validation" else: eval_loss, eval_acc, eval_mse = tr_loss, tr_acc, tr_mse score_src = "training" results.append( { "method": method_key, "seed": seed, "score_source": score_src, "train_loss": tr_loss, "train_accuracy": tr_acc, "train_mse": tr_mse, "eval_loss": eval_loss, "eval_accuracy": eval_acc, "eval_mse": eval_mse, "loss": eval_loss, "accuracy": eval_acc, "mse": eval_mse, "elapsed_time": elapsed_t, } ) print( f"{method_key:<18} {seed:<6} {eval_loss:<12.6f} {eval_acc:<12.6f} {eval_mse:<12.6f} {elapsed_t:<10.4f}" ) print("-" * 80) print("\nAggregate Summary (Mean across seeds):") print("=" * 80) if has_val: print( f"{'Method':<18} {'Mean Val Loss':<14} {'Mean Val Acc':<14} {'Mean Val MSE':<14} {'Mean Time (s)':<12}" ) else: print( f"{'Method':<18} {'Mean Loss':<12} {'Mean Acc':<12} {'Mean MSE':<12} {'Mean Time (s)':<12}" ) print("-" * 80) summary_list: list[dict[str, Any]] = [] for method_key in methods_to_test: method_runs = [r for r in results if r["method"] == method_key] if not method_runs: continue n_runs = len(method_runs) mean_tr_loss = sum(r["train_loss"] for r in method_runs) / n_runs mean_tr_acc = sum(r["train_accuracy"] for r in method_runs) / n_runs mean_tr_mse = sum(r["train_mse"] for r in method_runs) / n_runs mean_eval_loss = sum(r["eval_loss"] for r in method_runs) / n_runs mean_eval_acc = sum(r["eval_accuracy"] for r in method_runs) / n_runs mean_eval_mse = sum(r["eval_mse"] for r in method_runs) / n_runs mean_time = sum(r["elapsed_time"] for r in method_runs) / n_runs score_src = method_runs[0]["score_source"] summary_entry = { "method": method_key, "title": available_m_plugins[method_key].title, "source": available_m_plugins[method_key].source, "score_source": score_src, "mean_train_loss": mean_tr_loss, "mean_train_accuracy": mean_tr_acc, "mean_train_mse": mean_tr_mse, "mean_eval_loss": mean_eval_loss, "mean_eval_accuracy": mean_eval_acc, "mean_eval_mse": mean_eval_mse, "mean_loss": mean_eval_loss, "mean_accuracy": mean_eval_acc, "mean_mse": mean_eval_mse, "mean_elapsed_time": mean_time, "runs": n_runs, } summary_list.append(summary_entry) if has_val: print( f"{method_key:<18} {mean_eval_loss:<14.6f} {mean_eval_acc:<14.6f} {mean_eval_mse:<14.6f} {mean_time:<12.4f}" ) else: print( f"{method_key:<18} {mean_eval_loss:<12.6f} {mean_eval_acc:<12.6f} {mean_eval_mse:<12.6f} {mean_time:<12.4f}" ) print("=" * 80) if args.json_path: payload = { "dataset": args.dataset, "score_source": "validation" if has_val else "training", "selectors": { "initialization": args.initialization, "evaluation": eval_stage, "convergence": args.convergence, "refinement": refine_stage, }, "parameters": { "n_particles": args.n_particles, "epochs": args.epochs, "batch_size": args.batch_size, "fitness_size": fitness_size, "refinement_epochs": refinement_epochs, "refinement_lr": args.refinement_lr, "seeds": args.seeds, }, "results": results, "summary": summary_list, } with open(args.json_path, "w", encoding="utf-8") as f: json.dump(payload, f, indent=2) print(f"\nSaved comparison results JSON to: {args.json_path}") if __name__ == "__main__": main()