| |
| """Visualize A1 fit outputs for target-mask evaluation.""" |
|
|
| 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 |
|
|
| from a1_pipeline.io_utils import ensure_directory |
|
|
|
|
| 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="Visualize A1 core ROI fit results") |
| parser.add_argument( |
| "--fit-output-dir", |
| type=str, |
| default="/home/mohith/ds005345/outputs/a1_bootstrap/fit_results/Qwen_Qwen3-0.6B", |
| help="Directory containing run_a1_fit.py outputs", |
| ) |
| parser.add_argument( |
| "--output-dir", |
| type=str, |
| default=None, |
| help="Plot output directory (default: <fit-output-dir>/plots)", |
| ) |
| parser.add_argument( |
| "--metric", |
| type=str, |
| default="mean_corr", |
| choices=[ |
| "mean_corr", |
| "mean_r2", |
| "mean_2v2_accuracy", |
| "2v2", |
| "2v2_accuracy", |
| "two_v_two_accuracy", |
| ], |
| help="Metric used for plots", |
| ) |
| return parser |
|
|
|
|
| def _check_required_files(fit_output_dir: Path) -> tuple[Path, Path]: |
| layer_summary_path = fit_output_dir / "core_roi_layer_summary.csv" |
| best_summary_path = fit_output_dir / "core_roi_best_layer_summary.csv" |
|
|
| missing = [path for path in [layer_summary_path, best_summary_path] if not path.exists()] |
| if missing: |
| raise FileNotFoundError("Missing fit summary files: " + ", ".join(str(path) for path in missing)) |
|
|
| return layer_summary_path, best_summary_path |
|
|
|
|
| def _plot_protocol_heatmap( |
| layer_summary_df: pd.DataFrame, |
| protocol: str, |
| metric: str, |
| output_path: Path, |
| ) -> None: |
| protocol_df = layer_summary_df[layer_summary_df["protocol"] == protocol].copy() |
| if protocol_df.empty: |
| return |
|
|
| pivot_df = protocol_df.pivot(index="roi_name", columns="layer_idx", values=metric) |
| pivot_df = pivot_df.sort_index() |
|
|
| plt.figure(figsize=(max(8, 0.4 * len(pivot_df.columns)), 4.8)) |
| sns.heatmap(pivot_df, cmap="viridis", annot=False) |
| plt.title(f"{protocol} {metric} by Target Mask and Layer") |
| plt.xlabel("Layer") |
| plt.ylabel("Target Mask") |
| plt.tight_layout() |
| plt.savefig(output_path, dpi=180) |
| plt.close() |
|
|
|
|
| def _plot_best_layer_bar( |
| best_df: pd.DataFrame, |
| metric: str, |
| output_path: Path, |
| ) -> None: |
| if best_df.empty: |
| return |
|
|
| chart_df = best_df.copy() |
| chart_df["label"] = chart_df["roi_name"] + "\nL" + chart_df["layer_idx"].astype(int).astype(str) |
|
|
| plt.figure(figsize=(11, 5)) |
| sns.barplot(data=chart_df, x="label", y=metric, hue="protocol") |
| plt.title(f"Best Layer per Target Mask ({metric})") |
| plt.xlabel("ROI and selected layer") |
| plt.ylabel(metric) |
| plt.xticks(rotation=0) |
| plt.tight_layout() |
| plt.savefig(output_path, dpi=180) |
| plt.close() |
|
|
|
|
| def _plot_protocol_layer_curve( |
| layer_summary_df: pd.DataFrame, |
| metric: str, |
| output_path: Path, |
| ) -> None: |
| if layer_summary_df.empty: |
| return |
|
|
| curve_df = ( |
| layer_summary_df.groupby(["protocol", "layer_idx"], as_index=False)[metric] |
| .mean() |
| .sort_values(["protocol", "layer_idx"]) |
| ) |
|
|
| plt.figure(figsize=(10, 4.8)) |
| sns.lineplot(data=curve_df, x="layer_idx", y=metric, hue="protocol", marker="o") |
| plt.title(f"Average Target-Mask {metric} by Layer") |
| plt.xlabel("Layer") |
| plt.ylabel(metric) |
| plt.tight_layout() |
| plt.savefig(output_path, dpi=180) |
| plt.close() |
|
|
|
|
| def main() -> None: |
| parser = _build_parser() |
| args = parser.parse_args() |
|
|
| fit_output_dir = Path(args.fit_output_dir).resolve() |
| layer_summary_path, best_summary_path = _check_required_files(fit_output_dir=fit_output_dir) |
|
|
| plot_output_dir = Path(args.output_dir).resolve() if args.output_dir else fit_output_dir / "plots" |
| ensure_directory(plot_output_dir) |
|
|
| layer_summary_df = pd.read_csv(layer_summary_path) |
| best_summary_df = pd.read_csv(best_summary_path) |
|
|
| metric = _resolve_metric_column(str(args.metric)) |
| if metric not in layer_summary_df.columns: |
| raise ValueError( |
| f"Metric column '{metric}' not found in {layer_summary_path}. " |
| f"Available columns: {sorted(layer_summary_df.columns.tolist())}" |
| ) |
| if metric not in best_summary_df.columns: |
| raise ValueError( |
| f"Metric column '{metric}' not found in {best_summary_path}. " |
| f"Available columns: {sorted(best_summary_df.columns.tolist())}" |
| ) |
|
|
| for protocol in sorted(set(layer_summary_df["protocol"].tolist())): |
| heatmap_path = plot_output_dir / f"{protocol}_{metric}_heatmap.png" |
| _plot_protocol_heatmap( |
| layer_summary_df=layer_summary_df, |
| protocol=protocol, |
| metric=metric, |
| output_path=heatmap_path, |
| ) |
|
|
| _plot_best_layer_bar( |
| best_df=best_summary_df, |
| metric=metric, |
| output_path=plot_output_dir / f"best_layer_{metric}_bar.png", |
| ) |
|
|
| _plot_protocol_layer_curve( |
| layer_summary_df=layer_summary_df, |
| metric=metric, |
| output_path=plot_output_dir / f"protocol_layer_curve_{metric}.png", |
| ) |
|
|
| print("=" * 72) |
| print("A1 visualization complete") |
| print(f"Fit output directory: {fit_output_dir}") |
| print(f"Plot output directory: {plot_output_dir}") |
| print(f"Metric: {metric}") |
| print("=" * 72) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
