worldmodel/JEPA/lejepa-identifiability/experiments/analysis/plot_gennorm.py

"""
Generalized normal sweep across 4 mixings.

Emits three figures:
  1. fig_gennorm.pdf       — 4-panel R^2(h -> z) vs alpha for SIGReg/VICReg/InfoNCE
  2. fig_gennorm_orth.pdf  — 4-panel orthogonality error vs alpha (unconstrained ylim
                             to show InfoNCE excursions off the chart)
  3. fig_gennorm_main.pdf  — single-panel spiral-only headline figure for main text,
                             matching the Fig.~4b style of the paper

Usage:
    python analysis/plot_gennorm.py --results_dir results/gennorm/ --out figures/
"""
import argparse, glob, json, os, re
import numpy as np
import matplotlib.pyplot as plt
from collections import defaultdict

MIXINGS = [("spiral", "Spiral"), ("banana", "Banana"),
           ("sinusoid", "Sinusoid"), ("nvp", "NVP")]

MODES = ("lejepa", "whiten", "infonce")
COLORS = {"lejepa": "#d62728", "whiten": "#1f77b4", "infonce": "#2ca02c"}
LABELS = {"lejepa": "SIGReg", "whiten": "VICReg", "infonce": "InfoNCE"}

YLABELS = {
    "r2_hz_grid":               r"Linear identifiability  $R^2(h \to z)$",
    "r2_hz":                    r"Linear identifiability  $R^2(h \to z)$",
    "orth_err_normalized_grid": r"Orthogonality error  $\|\hat Q^\top \hat Q - I\|_F / \sqrt{n}$",
    "orth_err_normalized":      r"Orthogonality error  $\|\hat Q^\top \hat Q - I\|_F / \sqrt{n}$",
}

# Sensible y-limits per metric. R^2 is bounded in [0,1] so we clip there.
# Orthogonality error is unbounded above (Whitening/InfoNCE off-Gaussian can spike
# into the tens), so we use log scale and let matplotlib autoscale.
YLIMS = {
    "r2_hz_grid":               (-0.05, 1.05),
    "r2_hz":                    (-0.05, 1.05),
    "orth_err_normalized_grid": None,  # autoscale; log scale (see YSCALES) handles outliers
    "orth_err_normalized":      None,
}

# Y-axis scale per metric. Linear by default; log for orth error to compress
# off-Gaussian excursions while still showing structure near zero.
YSCALES = {
    "r2_hz_grid":               "linear",
    "r2_hz":                    "linear",
    "orth_err_normalized_grid": "log",
    "orth_err_normalized":      "log",
}


# ──────────────────────────────────────────────────────────────────────────
# Data loading
# ──────────────────────────────────────────────────────────────────────────
def load_groups(results_dir, metric):
    """groups[(mixing, mode, alpha)] -> list of seed values for `metric`."""
    groups = defaultdict(list)
    for path in sorted(glob.glob(os.path.join(results_dir, "*.json"))):
        with open(path) as f:
            r = json.load(f)
        if r.get("experiment") != "gennorm":
            continue
        alpha = r.get("source_alpha")
        if alpha is None:
            m = re.search(r"alpha=([\d.]+)", r.get("run_name", ""))
            if m:
                alpha = float(m.group(1))
        if alpha is None or metric not in r:
            continue
        groups[(r["mixing"], r["mode"], alpha)].append(r[metric])
    return groups


def curve(groups, mixing, mode):
    alphas = sorted({a for (mx, m, a) in groups if mx == mixing and m == mode})
    mu = np.array([np.mean(groups[(mixing, mode, a)]) for a in alphas])
    sd = np.array([np.std (groups[(mixing, mode, a)]) for a in alphas])
    return np.array(alphas), mu, sd


# ──────────────────────────────────────────────────────────────────────────
# Figure 1 & 2: 4-panel grids (one per metric)
# ──────────────────────────────────────────────────────────────────────────
def plot_grid(groups, metric, out_path):
    fig, axes = plt.subplots(1, 4, figsize=(13, 3.0), sharey=True)
    use_log = YSCALES[metric] == "log"

    for ax, (mix_key, mix_name) in zip(axes, MIXINGS):
        for mode in MODES:
            alphas, mu, sd = curve(groups, mix_key, mode)
            if len(alphas) == 0:
                continue
            ax.plot(alphas, mu, marker="o", ms=5, lw=1.8,
                    color=COLORS[mode], label=LABELS[mode], zorder=3)
            # On log axes, clip the lower edge of the band away from zero
            # so fill_between doesn't disappear / warn.
            lower = mu - sd
            if use_log:
                lower = np.maximum(lower, 1e-3)
            ax.fill_between(alphas, lower, mu + sd,
                            color=COLORS[mode], alpha=0.2, zorder=2)
        ax.set_xscale("log", base=2)
        if use_log:
            ax.set_yscale("log")
        ax.axvline(2.0, color="black", lw=0.7, ls="--", alpha=0.6, zorder=1)
        ax.set_xlabel(r"Source shape $\alpha$")
        ax.set_title(mix_name)
        if YLIMS[metric] is not None:
            ax.set_ylim(*YLIMS[metric])
        ax.grid(alpha=0.3, which="both" if use_log else "major")

    axes[0].set_ylabel(YLABELS[metric])
    axes[-1].legend(frameon=False, loc="best")

    fig.tight_layout()
    fig.savefig(out_path, bbox_inches="tight")
    print(f"Saved {out_path}")
    plt.close(fig)


# ──────────────────────────────────────────────────────────────────────────
# Figure 3: main-text single-panel headline (spiral mixing, all three methods)
# ──────────────────────────────────────────────────────────────────────────
def plot_main_panel(groups, out_path):
    """Single-panel spiral-only figure to sit next to the bound-verification panel
    in the main-text composite figure (Fig.~4b in the paper)."""
    FIGSIZE  = 0.8 * np.array((4.0, 3.0))
    LW       = 2.0
    MS       = 6
    FONTSIZE = 11

    rc_saved = plt.rcParams.copy()
    plt.rcParams.update({
        "font.size":         FONTSIZE,
        "axes.labelsize":    FONTSIZE,
        "xtick.labelsize":   FONTSIZE - 1,
        "ytick.labelsize":   FONTSIZE - 1,
        "legend.fontsize":   FONTSIZE - 1,
        "axes.spines.top":   False,
        "axes.spines.right": False,
    })

    fig, ax = plt.subplots(figsize=FIGSIZE)

    for mode in MODES:
        alphas, mu, sd = curve(groups, "spiral", mode)
        if len(alphas) == 0:
            continue
        ax.plot(alphas, mu, marker="o", ms=MS, lw=LW,
                color=COLORS[mode], label=LABELS[mode], zorder=3)
        ax.fill_between(alphas, mu - sd, mu + sd,
                        color=COLORS[mode], alpha=0.2, zorder=2)

    # Reference lines for canonical distributions
    ax.axvline(1.0, 0, 0.95, color="gray",  lw=0.8, ls=":",  alpha=0.7, zorder=1)
    ax.axvline(2.0, 0, 0.90, color="black", lw=0.8, ls="--", alpha=0.7, zorder=1)
    ax.text(1.0, 1.04, "Laplace", ha="center", va="bottom",
            fontsize=FONTSIZE - 1, color="gray")

    ax.set_xscale("log", base=2)
    ax.set_ylabel(r"Linearity")
    ax.set_ylim(-0.05, 1.12)
    ax.grid(alpha=0.3, which="both")
    ax.legend(frameon=False, loc="lower right")
    ax.set_xticks([2**(-2), 2, 16])
    ax.set_xticklabels([r"$\leftarrow$ sparse", "Gaussian", r"uniform $\rightarrow$"])

    fig.tight_layout()
    fig.savefig(out_path, bbox_inches="tight", dpi=500)
    print(f"Saved {out_path}")
    plt.close(fig)
    plt.rcParams.update(rc_saved)


# ──────────────────────────────────────────────────────────────────────────
# Driver
# ──────────────────────────────────────────────────────────────────────────
def main():
    p = argparse.ArgumentParser()
    p.add_argument("--results_dir", default="results/gennorm/")
    p.add_argument("--out",         default="figures/")
    args = p.parse_args()
    os.makedirs(args.out, exist_ok=True)

    # 4-panel R^2 grid (appendix)
    groups_r2 = load_groups(args.results_dir, "r2_hz_grid")
    plot_grid(groups_r2, "r2_hz_grid",
              os.path.join(args.out, "fig_gennorm.pdf"))

    # 4-panel orthogonality grid (appendix, autoscaled)
    groups_orth = load_groups(args.results_dir, "orth_err_normalized_grid")
    plot_grid(groups_orth, "orth_err_normalized_grid",
              os.path.join(args.out, "fig_gennorm_orth.pdf"))

    # Single-panel main-text headline (spiral, all three methods)
    plot_main_panel(groups_r2,
                    os.path.join(args.out, "fig_gennorm_main.pdf"))


if __name__ == "__main__":
    main()