| |
| """Compare A1 fit results across multiple model slugs. |
| |
| Reads ``core_roi_layer_summary.csv`` and ``core_roi_best_layer_summary.csv`` |
| from each ``--fit-dir`` and renders side-by-side plots tagged by model slug. |
| |
| Outputs go to a dedicated comparison directory so the per-model visualize |
| step (run_a1_visualize.py) is never overwritten. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| from pathlib import Path |
|
|
| import matplotlib |
|
|
| matplotlib.use("Agg") |
|
|
| import matplotlib.pyplot as plt |
| import pandas as pd |
| import seaborn as sns |
|
|
|
|
| METRIC_ALIAS_MAP: dict[str, str] = { |
| "2v2": "mean_2v2_accuracy", |
| "2v2_accuracy": "mean_2v2_accuracy", |
| "two_v_two_accuracy": "mean_2v2_accuracy", |
| } |
|
|
|
|
| def _resolve_metric_column(metric: str) -> str: |
| token = str(metric).strip().lower() |
| return METRIC_ALIAS_MAP.get(token, str(metric)) |
|
|
|
|
| def _build_parser() -> argparse.ArgumentParser: |
| parser = argparse.ArgumentParser(description="Compare A1 fits across models") |
| parser.add_argument( |
| "--fit-dir", |
| action="append", |
| required=True, |
| help=( |
| "Path to a fit_results/<model-slug> directory. Repeat for each model. " |
| "Optionally prefix with a label, e.g. 'qwen3=outputs/.../fit_results/Qwen_Qwen3-0.6B'." |
| ), |
| ) |
| parser.add_argument( |
| "--output-dir", |
| type=str, |
| required=True, |
| help="Where to write comparison plots and merged CSVs", |
| ) |
| parser.add_argument( |
| "--metric", |
| type=str, |
| default="mean_corr", |
| choices=[ |
| "mean_corr", |
| "mean_r2", |
| "mean_2v2_accuracy", |
| "2v2", |
| "2v2_accuracy", |
| "two_v_two_accuracy", |
| ], |
| ) |
| parser.add_argument( |
| "--title-suffix", |
| type=str, |
| default="", |
| help="Optional title suffix to add to comparison plots", |
| ) |
| return parser |
|
|
|
|
| def _split_label(entry: str) -> tuple[str, Path]: |
| if "=" in entry: |
| label, path = entry.split("=", 1) |
| return label.strip() or Path(path).name, Path(path).expanduser().resolve() |
| p = Path(entry).expanduser().resolve() |
| return p.name, p |
|
|
|
|
| def _load_fit_summaries( |
| entries: list[str], |
| ) -> tuple[pd.DataFrame, pd.DataFrame]: |
| layer_frames: list[pd.DataFrame] = [] |
| best_frames: list[pd.DataFrame] = [] |
| for entry in entries: |
| label, fit_dir = _split_label(entry) |
| layer_path = fit_dir / "core_roi_layer_summary.csv" |
| best_path = fit_dir / "core_roi_best_layer_summary.csv" |
| missing = [p for p in (layer_path, best_path) if not p.exists()] |
| if missing: |
| raise FileNotFoundError( |
| "Missing fit summaries for " |
| f"{label} ({fit_dir}): {[str(p) for p in missing]}" |
| ) |
|
|
| layer_df = pd.read_csv(layer_path) |
| layer_df["model_label"] = label |
| layer_df["fit_dir"] = str(fit_dir) |
| layer_frames.append(layer_df) |
|
|
| best_df = pd.read_csv(best_path) |
| best_df["model_label"] = label |
| best_df["fit_dir"] = str(fit_dir) |
| best_frames.append(best_df) |
|
|
| return ( |
| pd.concat(layer_frames, ignore_index=True), |
| pd.concat(best_frames, ignore_index=True), |
| ) |
|
|
|
|
| def _plot_layer_curve( |
| layer_df: pd.DataFrame, |
| metric: str, |
| output_path: Path, |
| title_suffix: str, |
| ) -> None: |
| if layer_df.empty: |
| return |
| curve_df = ( |
| layer_df.groupby(["model_label", "protocol", "layer_idx"], as_index=False)[metric] |
| .mean() |
| .sort_values(["model_label", "protocol", "layer_idx"]) |
| ) |
|
|
| protocols = sorted(curve_df["protocol"].unique()) |
| fig, axes = plt.subplots( |
| 1, len(protocols), figsize=(6 * len(protocols), 4.5), sharey=True, squeeze=False |
| ) |
| for ax, protocol in zip(axes[0], protocols): |
| sub = curve_df[curve_df["protocol"] == protocol] |
| sns.lineplot(data=sub, x="layer_idx", y=metric, hue="model_label", marker="o", ax=ax) |
| ax.set_title(f"Protocol {protocol}") |
| ax.set_xlabel("Layer") |
| ax.set_ylabel(metric) |
| suffix = f" — {title_suffix}" if title_suffix else "" |
| fig.suptitle(f"Avg target-mask {metric} by layer across models{suffix}") |
| fig.tight_layout() |
| fig.savefig(output_path, dpi=180) |
| plt.close(fig) |
|
|
|
|
| def _plot_best_layer_bar( |
| best_df: pd.DataFrame, |
| metric: str, |
| output_path: Path, |
| title_suffix: str, |
| ) -> None: |
| if best_df.empty: |
| return |
| chart_df = best_df.copy() |
| chart_df["roi_label"] = chart_df["roi_name"].astype(str) |
|
|
| protocols = sorted(chart_df["protocol"].unique()) |
| fig, axes = plt.subplots( |
| len(protocols), 1, figsize=(12, 4.5 * len(protocols)), squeeze=False |
| ) |
| for ax, protocol in zip(axes[:, 0], protocols): |
| sub = chart_df[chart_df["protocol"] == protocol] |
| sns.barplot(data=sub, x="roi_label", y=metric, hue="model_label", ax=ax) |
| ax.set_title(f"Best-layer {metric} per ROI — protocol {protocol}") |
| ax.set_xlabel("ROI") |
| ax.set_ylabel(metric) |
| ax.tick_params(axis="x", rotation=20) |
| suffix = f" — {title_suffix}" if title_suffix else "" |
| fig.suptitle(f"Best layer {metric} per target mask across models{suffix}") |
| fig.tight_layout() |
| fig.savefig(output_path, dpi=180) |
| plt.close(fig) |
|
|
|
|
| def _plot_overall_summary( |
| best_df: pd.DataFrame, |
| metric: str, |
| output_path: Path, |
| title_suffix: str, |
| ) -> None: |
| if best_df.empty: |
| return |
| summary_df = ( |
| best_df.groupby(["model_label", "protocol"], as_index=False)[metric] |
| .mean() |
| .sort_values(["protocol", "model_label"]) |
| ) |
| fig, ax = plt.subplots(figsize=(8, 4.8)) |
| sns.barplot(data=summary_df, x="protocol", y=metric, hue="model_label", ax=ax) |
| ax.set_title( |
| f"Mean best-layer {metric} per protocol{(' — ' + title_suffix) if title_suffix else ''}" |
| ) |
| ax.set_ylabel(metric) |
| ax.set_xlabel("Protocol") |
| fig.tight_layout() |
| fig.savefig(output_path, dpi=180) |
| plt.close(fig) |
|
|
|
|
| def main() -> None: |
| args = _build_parser().parse_args() |
| output_dir = Path(args.output_dir).expanduser().resolve() |
| output_dir.mkdir(parents=True, exist_ok=True) |
|
|
| layer_df, best_df = _load_fit_summaries(entries=list(args.fit_dir)) |
| metric = _resolve_metric_column(args.metric) |
| if metric not in layer_df.columns: |
| raise ValueError( |
| f"Metric column '{metric}' missing from layer summary. " |
| f"Available: {sorted(layer_df.columns.tolist())}" |
| ) |
| if metric not in best_df.columns: |
| raise ValueError( |
| f"Metric column '{metric}' missing from best summary. " |
| f"Available: {sorted(best_df.columns.tolist())}" |
| ) |
|
|
| layer_df.to_csv(output_dir / "merged_core_roi_layer_summary.csv", index=False) |
| best_df.to_csv(output_dir / "merged_core_roi_best_layer_summary.csv", index=False) |
|
|
| _plot_layer_curve( |
| layer_df=layer_df, |
| metric=metric, |
| output_path=output_dir / f"compare_layer_curve_{metric}.png", |
| title_suffix=args.title_suffix, |
| ) |
| _plot_best_layer_bar( |
| best_df=best_df, |
| metric=metric, |
| output_path=output_dir / f"compare_best_layer_bar_{metric}.png", |
| title_suffix=args.title_suffix, |
| ) |
| _plot_overall_summary( |
| best_df=best_df, |
| metric=metric, |
| output_path=output_dir / f"compare_overall_{metric}.png", |
| title_suffix=args.title_suffix, |
| ) |
|
|
| print("=" * 72) |
| print("A1 model comparison complete") |
| print(f"Models compared : {sorted(set(layer_df['model_label']))}") |
| print(f"Metric : {metric}") |
| print(f"Output directory: {output_dir}") |
| print("=" * 72) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
