"""Aggregate CEC-2022 sweep results and render all paper figures + tables. Reads every ``result.json`` under ``outputs/runs/cec2022/`` (or a user-supplied directory), pivots into per-(algorithm, function, dim) arrays of final best fitness, then writes: * ``outputs/tables/cec2022_mean_std_d{dim}.tex`` — Table 1 * ``outputs/tables/cec2022_wilcoxon_d{dim}.tex`` — Table 2 * ``outputs/figures/cec2022_friedman.json`` — Friedman omnibus per dim * ``outputs/figures/convergence_d{dim}_{func}.{pdf,png}`` — Figure 1 * ``outputs/figures/wilcoxon_heatmap_d{dim}.{pdf,png}`` — Figure 4 """ from __future__ import annotations import argparse import json from collections import defaultdict from pathlib import Path from typing import Any import numpy as np import pandas as pd from ahdcma.stats.tests import friedman_test, wilcoxon_pairwise from ahdcma.viz.ablation import plot_pairwise_pvalue_heatmap from ahdcma.viz.convergence import plot_convergence from ahdcma.viz.tables import mean_std_table, wilcoxon_pvalue_table def _load_results(root: Path) -> list[dict[str, Any]]: out: list[dict[str, Any]] = [] for jf in root.rglob("result.json"): try: data = json.loads(jf.read_text()) out.append(data) except json.JSONDecodeError: continue return out def _pivot( results: list[dict[str, Any]], ) -> dict[int, dict[tuple[str, str], list[float]]]: """Group final best_f by dim, then by (algo, func).""" by_dim: dict[int, dict[tuple[str, str], list[float]]] = defaultdict(lambda: defaultdict(list)) for r in results: by_dim[int(r["dim"])][r["algo"], r["func"]].append(float(r["best_f"])) return by_dim def _convergence_curves( results: list[dict[str, Any]], dim: int, func: str ) -> dict[str, list[list[float]]]: by_algo: dict[str, list[list[float]]] = defaultdict(list) for r in results: if int(r["dim"]) != dim or r["func"] != func: continue by_algo[r["algo"]].append(list(r["best_fitness_curve"])) return by_algo def render_all( sweep_root: str | Path = "outputs/runs/cec2022", fig_dir: str | Path = "outputs/figures", tab_dir: str | Path = "outputs/tables", *, headline_funcs: list[str] | None = None, ) -> None: fig_dir = Path(fig_dir) tab_dir = Path(tab_dir) fig_dir.mkdir(parents=True, exist_ok=True) tab_dir.mkdir(parents=True, exist_ok=True) results = _load_results(Path(sweep_root)) if not results: print(f"No result.json files under {sweep_root}; nothing to render.") return by_dim = _pivot(results) headline_funcs = headline_funcs or [ "F1_zakharov", "F5_levy", "F9_composition1", ] summary: dict[str, Any] = {"n_results": len(results), "by_dim": {}} for dim, table in by_dim.items(): algos = sorted({a for a, _ in table}) funcs = sorted({f for _, f in table}) # Mean / std table mean_df = pd.DataFrame(index=funcs, columns=algos, dtype=object) for f in funcs: for a in algos: arr = table.get((a, f), []) if arr: mean_df.at[f, a] = (float(np.mean(arr)), float(np.std(arr))) else: mean_df.at[f, a] = np.nan mean_std_table( mean_df, tab_dir / f"cec2022_mean_std_d{dim}.tex", caption=f"CEC-2022 mean $\\pm$ std final fitness ({dim}-D, 30 seeds).", label=f"tab:cec2022_mean_std_d{dim}", ) # Wilcoxon vs AHD-CMA per function if "ahdcma" in algos: p_rows: list[dict[str, Any]] = [] for f in funcs: bundle = {a: np.asarray(table[a, f]) for a in algos if (a, f) in table} if "ahdcma" not in bundle or len(bundle) < 2: continue seeds_min = min(len(v) for v in bundle.values()) bundle = {k: v[:seeds_min] for k, v in bundle.items()} try: df = wilcoxon_pairwise(bundle, baseline_name="ahdcma") except ValueError: continue row = {"function": f} for _, rr in df.iterrows(): row[rr["algo"]] = rr["p_value"] p_rows.append(row) if p_rows: p_df = pd.DataFrame(p_rows).set_index("function") wilcoxon_pvalue_table( p_df, tab_dir / f"cec2022_wilcoxon_d{dim}.tex", caption=f"Wilcoxon p-values vs AHD-CMA on CEC-2022 ({dim}-D).", label=f"tab:wilcoxon_d{dim}", ) # Per-algo p-value heatmap (function x algo, log10) p_for_heat = p_df.dropna(how="all") if not p_for_heat.empty: plot_pairwise_pvalue_heatmap( p_for_heat, fig_dir / f"wilcoxon_heatmap_d{dim}", title=f"Wilcoxon p-values vs AHD-CMA, CEC-2022 {dim}-D", ) # Friedman omnibus across functions # Stack into a matrix where each row is a function, each column an algo. from numpy.typing import NDArray friedman_in: dict[str, NDArray[np.float64]] = {} for a in algos: cols = [] for f in funcs: arr = table.get((a, f), []) if not arr: cols.append(np.nan) continue cols.append(float(np.median(arr))) friedman_in[a] = np.asarray(cols, dtype=np.float64) # Drop functions with any NaN to keep equal-length input valid = ~np.any(np.stack([np.isnan(v) for v in friedman_in.values()], axis=0), axis=0) friedman_in = {k: v[valid] for k, v in friedman_in.items()} if all(arr.size >= 3 for arr in friedman_in.values()): stat, p = friedman_test(friedman_in) else: stat, p = float("nan"), float("nan") summary["by_dim"][str(dim)] = { "n_funcs": len(funcs), "n_algos": len(algos), "friedman_statistic": stat, "friedman_p_value": p, } # Convergence figures for the headline functions for f in headline_funcs: curves = _convergence_curves(results, dim, f) if not curves: continue plot_convergence( curves, fig_dir / f"convergence_d{dim}_{f}", title=f"{f} ({dim}-D)", xlabel="Generation", ylabel="Best fitness", log_y=True, ) summary_path = fig_dir / "cec2022_summary.json" summary_path.write_text(json.dumps(summary, indent=2)) print(json.dumps(summary, indent=2)) def main() -> None: p = argparse.ArgumentParser() p.add_argument("--sweep-root", default="outputs/runs/cec2022") p.add_argument("--fig-dir", default="outputs/figures") p.add_argument("--tab-dir", default="outputs/tables") args = p.parse_args() render_all(args.sweep_root, args.fig_dir, args.tab_dir) if __name__ == "__main__": main()