path: root/scripts/trajectory_gap_distribution.py

#!/usr/bin/env python3
"""Match FA/BP trajectory gap distributions against a BP-path bridge predictor.

This is a first distributional bridge for realized FA loss gaps. For each
feedback seed B, it compares:

1. Empirical gap: train the actual FA trajectory and measure
   L(theta_T^FA) - L(theta_T^BP).
2. Predicted bridge gap: integrate the FA/BP gradient mismatch along the fixed
   BP trajectory,

       delta_T(B) = -eta sum_t [g_FA(theta_t^BP; B) - g_BP(theta_t^BP)],

   then evaluate L(theta_T^BP + delta_T(B)) - L(theta_T^BP).

This is not a final theorem. It is the first test of whether a trajectory-level
bridge can predict a full gap distribution, not merely correlations.
"""

from __future__ import annotations

import argparse
import csv
import json
from dataclasses import asdict, dataclass
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
from scipy import stats

import trajectory_mlp_fa as tm


@dataclass(frozen=True)
class RunConfig:
    input_dim: int
    hidden_widths: list[int]
    output_dim: int
    samples: int
    steps: int
    lr: float
    data_seed: int
    init_seed: int
    feedback_seed_start: int
    feedback_runs: int
    feedback_init: str
    feedback_scale: str
    noise_std: float
    outdir: str
    plot: bool


@dataclass(frozen=True)
class GapRow:
    feedback_seed: int
    empirical_gap: float
    predicted_bridge_gap: float
    empirical_final_loss: float
    predicted_bridge_loss: float
    bp_final_loss: float
    initial_hidden_gradient_cosine: float
    final_hidden_gradient_cosine: float
    initial_q_mean: float
    final_q_mean: float
    bridge_delta_norm: float
    empirical_delta_norm: float


@dataclass(frozen=True)
class MatchSummary:
    runs: int
    bp_final_loss: float
    empirical_mean: float
    predicted_mean: float
    empirical_std: float
    predicted_std: float
    empirical_q01: float
    predicted_q01: float
    empirical_q50: float
    predicted_q50: float
    empirical_q99: float
    predicted_q99: float
    ks_2sample_statistic: float
    ks_2sample_pvalue: float
    standardized_ks_statistic: float
    standardized_ks_pvalue: float
    moment_matched_ks_statistic: float
    moment_matched_ks_pvalue: float
    moment_matched_wasserstein_distance: float
    wasserstein_distance: float
    pearson_r: float
    pearson_p: float
    spearman_r: float
    spearman_p: float


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="FA/BP trajectory gap distribution matching."
    )
    parser.add_argument("--input-dim", type=int, default=16)
    parser.add_argument("--hidden-widths", type=int, nargs="+", default=[16, 16])
    parser.add_argument("--output-dim", type=int, default=4)
    parser.add_argument("--samples", type=int, default=256)
    parser.add_argument("--steps", type=int, default=120)
    parser.add_argument("--lr", type=float, default=0.02)
    parser.add_argument("--data-seed", type=int, default=20)
    parser.add_argument("--init-seed", type=int, default=30)
    parser.add_argument("--feedback-seed-start", type=int, default=5_000)
    parser.add_argument("--feedback-runs", type=int, default=500)
    parser.add_argument(
        "--feedback-init",
        choices=["gaussian", "rademacher"],
        default="gaussian",
    )
    parser.add_argument(
        "--feedback-scale",
        choices=["relu", "fan-in", "unit"],
        default="relu",
    )
    parser.add_argument("--noise-std", type=float, default=0.01)
    parser.add_argument(
        "--outdir",
        type=Path,
        default=Path("outputs/trajectory_gap_distribution"),
    )
    parser.add_argument("--plot", action="store_true")
    return parser.parse_args()


def parse_config(args: argparse.Namespace) -> RunConfig:
    return RunConfig(
        input_dim=args.input_dim,
        hidden_widths=args.hidden_widths,
        output_dim=args.output_dim,
        samples=args.samples,
        steps=args.steps,
        lr=args.lr,
        data_seed=args.data_seed,
        init_seed=args.init_seed,
        feedback_seed_start=args.feedback_seed_start,
        feedback_runs=args.feedback_runs,
        feedback_init=args.feedback_init,
        feedback_scale=args.feedback_scale,
        noise_std=args.noise_std,
        outdir=str(args.outdir),
        plot=args.plot,
    )


def tm_config(config: RunConfig) -> tm.RunConfig:
    return tm.RunConfig(
        input_dim=config.input_dim,
        hidden_widths=config.hidden_widths,
        output_dim=config.output_dim,
        samples=config.samples,
        steps=config.steps,
        lr=config.lr,
        eval_every=config.steps,
        data_seed=config.data_seed,
        init_seed=config.init_seed,
        feedback_seed_start=config.feedback_seed_start,
        feedback_runs=config.feedback_runs,
        feedback_init=config.feedback_init,
        feedback_scale=config.feedback_scale,
        noise_std=config.noise_std,
        outdir=config.outdir,
        plot=False,
    )


def validate_config(config: RunConfig) -> None:
    run_config = tm_config(config)
    tm.validate_config(run_config)


def zeros_like(weights: list[tm.Array]) -> list[tm.Array]:
    return [np.zeros_like(weight) for weight in weights]


def copy_weights(weights: list[tm.Array]) -> list[tm.Array]:
    return [weight.copy() for weight in weights]


def add_scaled(target: list[tm.Array], source: list[tm.Array], scale: float) -> None:
    for target_array, source_array in zip(target, source):
        target_array += scale * source_array


def add_weights(a: list[tm.Array], b: list[tm.Array]) -> list[tm.Array]:
    return [left + right for left, right in zip(a, b)]


def weights_delta(a: list[tm.Array], b: list[tm.Array]) -> list[tm.Array]:
    return [left - right for left, right in zip(a, b)]


def norm_weights(weights: list[tm.Array]) -> float:
    return float(np.linalg.norm(tm.flatten(weights)))


def train_bp_path(
    initial_weights: list[tm.Array], x: tm.Array, y: tm.Array, lr: float, steps: int
) -> tuple[list[list[tm.Array]], list[list[tm.Array]], list[tm.Array], float]:
    weights = copy_weights(initial_weights)
    path: list[list[tm.Array]] = []
    bp_grads: list[list[tm.Array]] = []

    for _step in range(steps):
        path.append(copy_weights(weights))
        grads, _ = tm.gradients(weights, x, y, feedback=None)
        bp_grads.append(copy_weights(grads))
        tm.sgd_step(weights, grads, lr)

    bp_final_loss = tm.mse_loss(tm.predict(weights, x), y)
    return path, bp_grads, weights, bp_final_loss


def bridge_delta(
    bp_path: list[list[tm.Array]],
    bp_grads: list[list[tm.Array]],
    feedback: list[tm.Array],
    x: tm.Array,
    y: tm.Array,
    lr: float,
) -> list[tm.Array]:
    delta = zeros_like(bp_path[0])
    for weights, bp_grad in zip(bp_path, bp_grads):
        fa_grad, _ = tm.gradients(weights, x, y, feedback=feedback)
        mismatch = [fa - bp for fa, bp in zip(fa_grad, bp_grad)]
        add_scaled(delta, mismatch, -lr)
    return delta


def run_one_feedback(
    config: RunConfig,
    run_config: tm.RunConfig,
    initial_weights: list[tm.Array],
    bp_path: list[list[tm.Array]],
    bp_grads: list[list[tm.Array]],
    bp_final_weights: list[tm.Array],
    bp_final_loss: float,
    x: tm.Array,
    y: tm.Array,
    feedback_seed: int,
) -> GapRow:
    widths = tm.layer_widths(run_config)
    feedback = tm.init_feedback(
        widths, feedback_seed, config.feedback_init, config.feedback_scale
    )

    delta = bridge_delta(bp_path, bp_grads, feedback, x, y, config.lr)
    bridge_weights = add_weights(bp_final_weights, delta)
    bridge_loss = tm.mse_loss(tm.predict(bridge_weights, x), y)
    predicted_gap = bridge_loss - bp_final_loss

    fa_final_weights, trajectory, _layer_metrics = tm.train_fa(
        initial_weights, feedback, feedback_seed, x, y, run_config
    )
    empirical_final_loss = trajectory[-1].loss
    empirical_gap = empirical_final_loss - bp_final_loss

    initial = trajectory[0]
    final = trajectory[-1]
    empirical_delta = weights_delta(fa_final_weights, bp_final_weights)

    return GapRow(
        feedback_seed=feedback_seed,
        empirical_gap=empirical_gap,
        predicted_bridge_gap=predicted_gap,
        empirical_final_loss=empirical_final_loss,
        predicted_bridge_loss=bridge_loss,
        bp_final_loss=bp_final_loss,
        initial_hidden_gradient_cosine=float(initial.hidden_gradient_cosine),
        final_hidden_gradient_cosine=float(final.hidden_gradient_cosine),
        initial_q_mean=float(initial.q_mean),
        final_q_mean=float(final.q_mean),
        bridge_delta_norm=norm_weights(delta),
        empirical_delta_norm=norm_weights(empirical_delta),
    )


def summarize(rows: list[GapRow], bp_final_loss: float) -> MatchSummary:
    empirical = np.array([row.empirical_gap for row in rows], dtype=np.float64)
    predicted = np.array([row.predicted_bridge_gap for row in rows], dtype=np.float64)
    ks = stats.ks_2samp(predicted, empirical)
    predicted_std = np.std(predicted, ddof=1)
    empirical_std = np.std(empirical, ddof=1)
    predicted_z = (predicted - np.mean(predicted)) / predicted_std
    empirical_z = (empirical - np.mean(empirical)) / empirical_std
    standardized_ks = stats.ks_2samp(predicted_z, empirical_z)
    moment_matched = np.mean(empirical) + predicted_z * empirical_std
    moment_matched_ks = stats.ks_2samp(moment_matched, empirical)
    pearson = stats.pearsonr(predicted, empirical)
    spearman = stats.spearmanr(predicted, empirical)
    return MatchSummary(
        runs=len(rows),
        bp_final_loss=bp_final_loss,
        empirical_mean=float(np.mean(empirical)),
        predicted_mean=float(np.mean(predicted)),
        empirical_std=float(empirical_std),
        predicted_std=float(predicted_std),
        empirical_q01=float(np.quantile(empirical, 0.01)),
        predicted_q01=float(np.quantile(predicted, 0.01)),
        empirical_q50=float(np.quantile(empirical, 0.50)),
        predicted_q50=float(np.quantile(predicted, 0.50)),
        empirical_q99=float(np.quantile(empirical, 0.99)),
        predicted_q99=float(np.quantile(predicted, 0.99)),
        ks_2sample_statistic=float(ks.statistic),
        ks_2sample_pvalue=float(ks.pvalue),
        standardized_ks_statistic=float(standardized_ks.statistic),
        standardized_ks_pvalue=float(standardized_ks.pvalue),
        moment_matched_ks_statistic=float(moment_matched_ks.statistic),
        moment_matched_ks_pvalue=float(moment_matched_ks.pvalue),
        moment_matched_wasserstein_distance=float(
            stats.wasserstein_distance(moment_matched, empirical)
        ),
        wasserstein_distance=float(stats.wasserstein_distance(predicted, empirical)),
        pearson_r=float(pearson.statistic),
        pearson_p=float(pearson.pvalue),
        spearman_r=float(spearman.statistic),
        spearman_p=float(spearman.pvalue),
    )


def write_csv(path: Path, rows: list[object]) -> None:
    if not rows:
        return
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", newline="") as handle:
        first = asdict(rows[0])  # type: ignore[arg-type]
        writer = csv.DictWriter(handle, fieldnames=list(first.keys()))
        writer.writeheader()
        for row in rows:
            writer.writerow(asdict(row))  # type: ignore[arg-type]


def write_outputs(
    config: RunConfig, rows: list[GapRow], summary: MatchSummary, outdir: Path
) -> None:
    outdir.mkdir(parents=True, exist_ok=True)
    write_csv(outdir / "gap_distribution.csv", rows)
    write_csv(outdir / "summary.csv", [summary])
    payload = {
        "config": asdict(config),
        "summary": asdict(summary),
        "runs": [asdict(row) for row in rows],
    }
    (outdir / "summary.json").write_text(
        json.dumps(payload, indent=2, sort_keys=True) + "\n"
    )


def save_plots(rows: list[GapRow], outdir: Path) -> list[Path]:
    outdir.mkdir(parents=True, exist_ok=True)
    paths: list[Path] = []
    empirical = np.array([row.empirical_gap for row in rows], dtype=np.float64)
    predicted = np.array([row.predicted_bridge_gap for row in rows], dtype=np.float64)
    predicted_z = (predicted - np.mean(predicted)) / np.std(predicted, ddof=1)
    empirical_z = (empirical - np.mean(empirical)) / np.std(empirical, ddof=1)
    moment_matched = np.mean(empirical) + predicted_z * np.std(empirical, ddof=1)

    hist_path = outdir / "gap_histogram_overlay.png"
    bins = np.histogram_bin_edges(np.concatenate([empirical, predicted]), bins=60)
    plt.figure(figsize=(7, 4.5))
    plt.hist(predicted, bins=bins, density=True, alpha=0.45, label="bridge predicted")
    plt.hist(empirical, bins=bins, density=True, alpha=0.45, label="actual FA")
    plt.xlabel("final loss gap to BP")
    plt.ylabel("density")
    plt.title("Predicted vs empirical FA/BP gap distribution")
    plt.legend()
    plt.tight_layout()
    plt.savefig(hist_path, dpi=180)
    plt.close()
    paths.append(hist_path)

    matched_hist_path = outdir / "gap_histogram_moment_matched.png"
    bins = np.histogram_bin_edges(np.concatenate([empirical, moment_matched]), bins=60)
    plt.figure(figsize=(7, 4.5))
    plt.hist(moment_matched, bins=bins, density=True, alpha=0.45, label="bridge moment-matched")
    plt.hist(empirical, bins=bins, density=True, alpha=0.45, label="actual FA")
    plt.xlabel("final loss gap to BP")
    plt.ylabel("density")
    plt.title("Moment-matched bridge vs empirical gap distribution")
    plt.legend()
    plt.tight_layout()
    plt.savefig(matched_hist_path, dpi=180)
    plt.close()
    paths.append(matched_hist_path)

    cdf_path = outdir / "gap_cdf_overlay.png"
    plt.figure(figsize=(7, 4.5))
    for values, label in [(predicted, "bridge predicted"), (empirical, "actual FA")]:
        sorted_values = np.sort(values)
        probs = (np.arange(1, len(values) + 1) - 0.5) / len(values)
        plt.plot(sorted_values, probs, label=label)
    plt.xlabel("final loss gap to BP")
    plt.ylabel("CDF")
    plt.title("Gap distribution CDF")
    plt.legend()
    plt.tight_layout()
    plt.savefig(cdf_path, dpi=180)
    plt.close()
    paths.append(cdf_path)

    qq_path = outdir / "gap_qq_plot.png"
    probs = (np.arange(1, len(empirical) + 1) - 0.5) / len(empirical)
    pred_q = np.quantile(predicted, probs)
    emp_q = np.quantile(empirical, probs)
    min_value = float(min(pred_q[0], emp_q[0]))
    max_value = float(max(pred_q[-1], emp_q[-1]))
    plt.figure(figsize=(5, 5))
    plt.scatter(pred_q, emp_q, s=8, alpha=0.35)
    plt.plot([min_value, max_value], [min_value, max_value], color="black", linewidth=1)
    plt.xlabel("bridge predicted quantile")
    plt.ylabel("actual FA quantile")
    plt.title("Gap distribution Q-Q")
    plt.tight_layout()
    plt.savefig(qq_path, dpi=180)
    plt.close()
    paths.append(qq_path)

    standardized_qq_path = outdir / "gap_standardized_qq_plot.png"
    pred_q = np.quantile(predicted_z, probs)
    emp_q = np.quantile(empirical_z, probs)
    min_value = float(min(pred_q[0], emp_q[0]))
    max_value = float(max(pred_q[-1], emp_q[-1]))
    plt.figure(figsize=(5, 5))
    plt.scatter(pred_q, emp_q, s=8, alpha=0.35)
    plt.plot([min_value, max_value], [min_value, max_value], color="black", linewidth=1)
    plt.xlabel("standardized bridge quantile")
    plt.ylabel("standardized actual FA quantile")
    plt.title("Standardized gap distribution Q-Q")
    plt.tight_layout()
    plt.savefig(standardized_qq_path, dpi=180)
    plt.close()
    paths.append(standardized_qq_path)

    scatter_path = outdir / "paired_gap_scatter.png"
    plt.figure(figsize=(5, 5))
    plt.scatter(predicted, empirical, s=12, alpha=0.45)
    min_value = float(min(np.min(predicted), np.min(empirical)))
    max_value = float(max(np.max(predicted), np.max(empirical)))
    plt.plot([min_value, max_value], [min_value, max_value], color="black", linewidth=1)
    plt.xlabel("bridge predicted gap")
    plt.ylabel("actual FA gap")
    plt.title("Paired feedback seeds")
    plt.tight_layout()
    plt.savefig(scatter_path, dpi=180)
    plt.close()
    paths.append(scatter_path)

    return paths


def main() -> None:
    args = parse_args()
    config = parse_config(args)
    validate_config(config)
    run_config = tm_config(config)
    widths = tm.layer_widths(run_config)
    x, y = tm.make_synthetic_regression(run_config)
    initial_weights = tm.init_weights(widths, config.init_seed)
    bp_path, bp_grads, bp_final_weights, bp_final_loss = train_bp_path(
        initial_weights, x, y, config.lr, config.steps
    )

    rows: list[GapRow] = []
    for run_index in range(config.feedback_runs):
        feedback_seed = config.feedback_seed_start + run_index
        row = run_one_feedback(
            config,
            run_config,
            initial_weights,
            bp_path,
            bp_grads,
            bp_final_weights,
            bp_final_loss,
            x,
            y,
            feedback_seed,
        )
        rows.append(row)
        if (run_index + 1) % max(1, config.feedback_runs // 10) == 0:
            print(
                f"completed {run_index + 1}/{config.feedback_runs}: "
                f"emp_gap={row.empirical_gap:.6g}, pred_gap={row.predicted_bridge_gap:.6g}"
            )

    summary = summarize(rows, bp_final_loss)
    outdir = Path(config.outdir)
    write_outputs(config, rows, summary, outdir)
    plot_paths = save_plots(rows, outdir) if config.plot else []

    print(f"widths: {widths}")
    print(f"bp_final_loss: {bp_final_loss:.8g}")
    print(
        "empirical_gap: "
        f"mean={summary.empirical_mean:.8g}, std={summary.empirical_std:.8g}"
    )
    print(
        "predicted_bridge_gap: "
        f"mean={summary.predicted_mean:.8g}, std={summary.predicted_std:.8g}"
    )
    print(
        "distribution_match: "
        f"KS={summary.ks_2sample_statistic:.8g}, "
        f"W1={summary.wasserstein_distance:.8g}, "
        f"Spearman={summary.spearman_r:.8g}"
    )
    print(
        "shape_match_after_standardization: "
        f"KS={summary.standardized_ks_statistic:.8g}, "
        f"p={summary.standardized_ks_pvalue:.8g}, "
        f"moment_matched_W1={summary.moment_matched_wasserstein_distance:.8g}"
    )
    print(f"gap_distribution: {outdir / 'gap_distribution.csv'}")
    print(f"summary: {outdir / 'summary.csv'}")
    for path in plot_paths:
        print(f"plot: {path}")


if __name__ == "__main__":
    main()