code/run_a1_visualize.py

#!/usr/bin/env python
"""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()