diff options
Diffstat (limited to 'scripts/plot_tangent_hierarchy_first_order_try.py')
| -rw-r--r-- | scripts/plot_tangent_hierarchy_first_order_try.py | 166 |
1 files changed, 166 insertions, 0 deletions
diff --git a/scripts/plot_tangent_hierarchy_first_order_try.py b/scripts/plot_tangent_hierarchy_first_order_try.py new file mode 100644 index 0000000..c74c8d2 --- /dev/null +++ b/scripts/plot_tangent_hierarchy_first_order_try.py @@ -0,0 +1,166 @@ +#!/usr/bin/env python3 +"""Plot first-order FA tangent-hierarchy predictor diagnostics.""" + +from __future__ import annotations + +import argparse +from pathlib import Path + +import matplotlib.pyplot as plt +import numpy as np +import pandas as pd + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument( + "--csv", + type=Path, + default=Path("outputs/fa_tangent_hierarchy_first_order_try/compressed_operator_rows.csv"), + ) + parser.add_argument( + "--outdir", + type=Path, + default=Path("outputs/fa_tangent_hierarchy_first_order_try"), + ) + return parser.parse_args() + + +def metrics(df: pd.DataFrame) -> pd.DataFrame: + rows = [] + for early, group in df.groupby("early_steps"): + for label, column in [ + ("fixed K0", "fixed_gap"), + ("linear velocity", "linear_gap"), + ("retangent", "retangent_gap"), + ("compressed", "compressed_gap"), + ]: + error = group[column] - group["empirical_gap"] + rows.append( + { + "early_steps": int(early), + "predictor": label, + "mae": float(error.abs().mean()), + "bias": float(error.mean()), + "rmse": float(np.sqrt(np.mean(np.square(error)))), + "corr": float(group[column].corr(group["empirical_gap"])), + } + ) + return pd.DataFrame(rows) + + +def plot_scatter(df: pd.DataFrame, outdir: Path) -> Path: + early_values = sorted(df["early_steps"].unique()) + fig, axes = plt.subplots(1, len(early_values) + 1, figsize=(5.0 * (len(early_values) + 1), 4.6), dpi=170) + + panels = [("fixed K0", "fixed_gap", df)] + for early in early_values: + panels.append((f"linear velocity s={early}", "linear_gap", df[df["early_steps"] == early])) + + values = [] + for _title, column, group in panels: + values.extend(group[column].tolist()) + values.extend(group["empirical_gap"].tolist()) + lo, hi = float(min(values)), float(max(values)) + pad = 0.04 * (hi - lo + 1e-12) + + for ax, (title, column, group) in zip(axes, panels): + ax.scatter(group[column], group["empirical_gap"], s=26, alpha=0.72, color="#1f5a9d") + ax.plot([lo - pad, hi + pad], [lo - pad, hi + pad], color="black", linewidth=1.0) + err = group[column] - group["empirical_gap"] + ax.set_title(f"{title}\nMAE={err.abs().mean():.4f}, bias={err.mean():+.4f}") + ax.set_xlabel("predicted gap") + ax.set_ylabel("empirical gap") + ax.grid(alpha=0.18) + ax.set_xlim(lo - pad, hi + pad) + ax.set_ylim(lo - pad, hi + pad) + + fig.suptitle("First-order tangent-hierarchy predictor", y=1.02) + fig.tight_layout() + path = outdir / "first_order_prediction_scatter_by_early_step.png" + fig.savefig(path, bbox_inches="tight") + plt.close(fig) + return path + + +def plot_metrics(metric_df: pd.DataFrame, outdir: Path) -> Path: + fig, axes = plt.subplots(1, 3, figsize=(13.5, 4.4), dpi=170) + predictors = ["fixed K0", "linear velocity", "retangent", "compressed"] + colors = { + "fixed K0": "#777777", + "linear velocity": "#1f5a9d", + "retangent": "#2b8a3e", + "compressed": "#8a4fb0", + } + + for ax, metric_name in zip(axes, ["mae", "bias", "corr"]): + for predictor in predictors: + sub = metric_df[metric_df["predictor"] == predictor].sort_values("early_steps") + ax.plot( + sub["early_steps"], + sub[metric_name], + marker="o", + linewidth=1.8, + color=colors[predictor], + label=predictor, + ) + if metric_name == "bias": + ax.axhline(0.0, color="black", linewidth=1.0) + ax.set_xlabel("early operator step s") + ax.set_title(metric_name.upper()) + ax.grid(alpha=0.18) + axes[0].set_ylabel("metric value") + axes[0].legend(fontsize=8) + fig.suptitle("Prediction metrics vs early operator horizon", y=1.02) + fig.tight_layout() + path = outdir / "first_order_error_metrics_by_early_step.png" + fig.savefig(path, bbox_inches="tight") + plt.close(fig) + return path + + +def plot_distributions(df: pd.DataFrame, outdir: Path) -> Path: + early_values = sorted(df["early_steps"].unique()) + fig, axes = plt.subplots(1, len(early_values), figsize=(5.0 * len(early_values), 4.4), dpi=170, squeeze=False) + for ax, early in zip(axes.ravel(), early_values): + group = df[df["early_steps"] == early] + bins = np.linspace( + min(group["empirical_gap"].min(), group["linear_gap"].min(), group["fixed_gap"].min()), + max(group["empirical_gap"].max(), group["linear_gap"].max(), group["fixed_gap"].max()), + 18, + ) + ax.hist(group["empirical_gap"], bins=bins, alpha=0.45, density=True, color="#c65f16", label="empirical") + ax.hist(group["fixed_gap"], bins=bins, alpha=0.32, density=True, color="#777777", label="fixed K0") + ax.hist(group["linear_gap"], bins=bins, alpha=0.36, density=True, color="#1f5a9d", label="linear velocity") + ax.set_title(f"gap distribution, s={early}") + ax.set_xlabel("FA/BP train gap") + ax.set_ylabel("density") + ax.legend(fontsize=8) + ax.grid(alpha=0.18) + fig.tight_layout() + path = outdir / "first_order_gap_distribution_overlay.png" + fig.savefig(path, bbox_inches="tight") + plt.close(fig) + return path + + +def main() -> None: + args = parse_args() + args.outdir.mkdir(parents=True, exist_ok=True) + df = pd.read_csv(args.csv) + metric_df = metrics(df) + metric_path = args.outdir / "first_order_metrics.csv" + metric_df.to_csv(metric_path, index=False) + paths = [ + plot_scatter(df, args.outdir), + plot_metrics(metric_df, args.outdir), + plot_distributions(df, args.outdir), + ] + print(f"metrics: {metric_path}") + for path in paths: + print(f"plot: {path}") + print(metric_df.to_string(index=False)) + + +if __name__ == "__main__": + main() |