1 files changed, 285 insertions, 0 deletions
diff --git a/scripts/static_alignment_beta.py b/scripts/static_alignment_beta.py
new file mode 100755
index 0000000..a53e6fa
--- /dev/null
+++ b/scripts/static_alignment_beta.py
@@ -0,0 +1,285 @@
+#!/usr/bin/env python3
+"""Validate the beta law for random feedback alignment.
+
+For independent isotropic matrix directions A, B in R^{rows x cols}, the
+squared Frobenius cosine
+
+    Q = <A, B>_F^2 / (||A||_F^2 ||B||_F^2)
+
+has distribution Beta(1/2, (D - 1)/2), where D = rows * cols.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+from dataclasses import asdict, dataclass
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+import numpy as np
+from scipy import special, stats
+
+
+@dataclass(frozen=True)
+class RunConfig:
+    rows: int
+    cols: int
+    samples: int
+    batch_size: int
+    seed: int
+    dist_a: str
+    dist_b: str
+    thresholds: list[float]
+    outdir: str
+    plot: bool
+
+
+@dataclass(frozen=True)
+class Summary:
+    dimension: int
+    beta_alpha: float
+    beta_beta: float
+    empirical_mean: float
+    theoretical_mean: float
+    empirical_var: float
+    theoretical_var: float
+    ks_statistic: float
+    ks_pvalue: float
+    quantiles: dict[str, float]
+    tail_probabilities: dict[str, dict[str, float]]
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description="Monte Carlo validation of the FA static alignment beta law."
+    )
+    parser.add_argument("--rows", type=int, default=16, help="Matrix row count.")
+    parser.add_argument("--cols", type=int, default=16, help="Matrix column count.")
+    parser.add_argument("--samples", type=int, default=50_000, help="Number of pairs.")
+    parser.add_argument(
+        "--batch-size",
+        type=int,
+        default=10_000,
+        help="Number of pairs to sample per batch.",
+    )
+    parser.add_argument("--seed", type=int, default=0, help="Random seed.")
+    parser.add_argument(
+        "--dist-a",
+        choices=["gaussian", "sphere", "rademacher"],
+        default="gaussian",
+        help="Distribution for A.",
+    )
+    parser.add_argument(
+        "--dist-b",
+        choices=["gaussian", "sphere", "rademacher"],
+        default="gaussian",
+        help="Distribution for B.",
+    )
+    parser.add_argument(
+        "--thresholds",
+        type=float,
+        nargs="*",
+        default=None,
+        help="Tail thresholds q. Defaults to {1, 2, 5, 10}/D.",
+    )
+    parser.add_argument(
+        "--outdir",
+        type=Path,
+        default=Path("outputs/static_alignment_beta"),
+        help="Directory for summary and plots.",
+    )
+    parser.add_argument("--plot", action="store_true", help="Save diagnostic plots.")
+    return parser.parse_args()
+
+
+def sample_vectors(
+    rng: np.random.Generator, count: int, dim: int, distribution: str
+) -> np.ndarray:
+    if distribution == "gaussian":
+        return rng.standard_normal((count, dim))
+
+    if distribution == "sphere":
+        values = rng.standard_normal((count, dim))
+        norms = np.linalg.norm(values, axis=1, keepdims=True)
+        return values / np.maximum(norms, np.finfo(values.dtype).tiny)
+
+    if distribution == "rademacher":
+        return rng.choice(np.array([-1.0, 1.0]), size=(count, dim))
+
+    raise ValueError(f"Unknown distribution: {distribution}")
+
+
+def squared_cosines(config: RunConfig) -> np.ndarray:
+    dim = config.rows * config.cols
+    rng = np.random.default_rng(config.seed)
+    values = np.empty(config.samples, dtype=np.float64)
+
+    offset = 0
+    while offset < config.samples:
+        count = min(config.batch_size, config.samples - offset)
+        a = sample_vectors(rng, count, dim, config.dist_a)
+        b = sample_vectors(rng, count, dim, config.dist_b)
+
+        dot = np.einsum("ij,ij->i", a, b)
+        a_norm_sq = np.einsum("ij,ij->i", a, a)
+        b_norm_sq = np.einsum("ij,ij->i", b, b)
+        values[offset : offset + count] = (dot * dot) / (a_norm_sq * b_norm_sq)
+        offset += count
+
+    return values
+
+
+def summarize(q_values: np.ndarray, dimension: int, thresholds: list[float]) -> Summary:
+    alpha = 0.5
+    beta_param = (dimension - 1) / 2
+    beta_dist = stats.beta(alpha, beta_param)
+
+    quantile_levels = [0.01, 0.05, 0.10, 0.50, 0.90, 0.95, 0.99]
+    quantiles = {
+        f"{level:.2f}": float(np.quantile(q_values, level))
+        for level in quantile_levels
+    }
+
+    tail_probabilities: dict[str, dict[str, float]] = {}
+    for threshold in thresholds:
+        empirical_tail = float(np.mean(q_values >= threshold))
+        theoretical_tail = float(1.0 - special.betainc(alpha, beta_param, threshold))
+        tail_probabilities[f"{threshold:.12g}"] = {
+            "empirical": empirical_tail,
+            "theoretical": theoretical_tail,
+            "absolute_error": abs(empirical_tail - theoretical_tail),
+        }
+
+    ks = stats.kstest(q_values, beta_dist.cdf)
+
+    return Summary(
+        dimension=dimension,
+        beta_alpha=alpha,
+        beta_beta=beta_param,
+        empirical_mean=float(np.mean(q_values)),
+        theoretical_mean=float(beta_dist.mean()),
+        empirical_var=float(np.var(q_values, ddof=1)),
+        theoretical_var=float(beta_dist.var()),
+        ks_statistic=float(ks.statistic),
+        ks_pvalue=float(ks.pvalue),
+        quantiles=quantiles,
+        tail_probabilities=tail_probabilities,
+    )
+
+
+def write_summary(config: RunConfig, summary: Summary, outdir: Path) -> Path:
+    outdir.mkdir(parents=True, exist_ok=True)
+    path = outdir / "summary.json"
+    payload = {
+        "config": asdict(config),
+        "summary": asdict(summary),
+    }
+    path.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n")
+    return path
+
+
+def save_plots(q_values: np.ndarray, dimension: int, outdir: Path) -> list[Path]:
+    outdir.mkdir(parents=True, exist_ok=True)
+    alpha = 0.5
+    beta_param = (dimension - 1) / 2
+    beta_dist = stats.beta(alpha, beta_param)
+
+    paths: list[Path] = []
+
+    hist_path = outdir / "histogram_beta_overlay.png"
+    x_min = float(max(beta_dist.ppf(1e-5), np.finfo(float).tiny))
+    x_max = float(max(np.quantile(q_values, 0.999), beta_dist.ppf(0.999)))
+    xs = np.linspace(x_min, x_max, 500)
+    beta_pdf = np.exp(np.clip(beta_dist.logpdf(xs), -745, 80))
+    plt.figure(figsize=(7, 4.5))
+    plt.hist(q_values, bins=80, density=True, alpha=0.45, label="empirical")
+    plt.plot(xs, beta_pdf, color="black", linewidth=2, label="beta law")
+    plt.xlabel("Q = squared Frobenius cosine")
+    plt.ylabel("density")
+    plt.title(f"Static alignment distribution, D={dimension}")
+    plt.legend()
+    plt.tight_layout()
+    plt.savefig(hist_path, dpi=180)
+    plt.close()
+    paths.append(hist_path)
+
+    qq_path = outdir / "qq_plot.png"
+    probs = (np.arange(1, len(q_values) + 1) - 0.5) / len(q_values)
+    empirical = np.sort(q_values)
+    theoretical = beta_dist.ppf(probs)
+    max_value = float(max(empirical[-1], theoretical[-1]))
+    plt.figure(figsize=(5, 5))
+    plt.scatter(theoretical, empirical, s=4, alpha=0.35)
+    plt.plot([0, max_value], [0, max_value], color="black", linewidth=1)
+    plt.xlabel("theoretical beta quantile")
+    plt.ylabel("empirical quantile")
+    plt.title("Q-Q plot")
+    plt.tight_layout()
+    plt.savefig(qq_path, dpi=180)
+    plt.close()
+    paths.append(qq_path)
+
+    return paths
+
+
+def main() -> None:
+    args = parse_args()
+    dimension = args.rows * args.cols
+    if dimension < 2:
+        raise ValueError("rows * cols must be at least 2 for the beta law.")
+    if args.samples < 1:
+        raise ValueError("--samples must be positive.")
+    if args.batch_size < 1:
+        raise ValueError("--batch-size must be positive.")
+
+    thresholds = args.thresholds
+    if thresholds is None:
+        thresholds = [1 / dimension, 2 / dimension, 5 / dimension, 10 / dimension]
+    thresholds = [q for q in thresholds if 0 <= q <= 1]
+
+    config = RunConfig(
+        rows=args.rows,
+        cols=args.cols,
+        samples=args.samples,
+        batch_size=args.batch_size,
+        seed=args.seed,
+        dist_a=args.dist_a,
+        dist_b=args.dist_b,
+        thresholds=thresholds,
+        outdir=str(args.outdir),
+        plot=args.plot,
+    )
+
+    q_values = squared_cosines(config)
+    summary = summarize(q_values, dimension, thresholds)
+    summary_path = write_summary(config, summary, args.outdir)
+
+    plot_paths: list[Path] = []
+    if args.plot:
+        plot_paths = save_plots(q_values, dimension, args.outdir)
+
+    print(f"dimension: {dimension}")
+    print(f"samples: {args.samples}")
+    print(
+        "mean: "
+        f"empirical={summary.empirical_mean:.8g}, "
+        f"theoretical={summary.theoretical_mean:.8g}"
+    )
+    print(
+        "variance: "
+        f"empirical={summary.empirical_var:.8g}, "
+        f"theoretical={summary.theoretical_var:.8g}"
+    )
+    print(
+        "KS: "
+        f"statistic={summary.ks_statistic:.6g}, "
+        f"pvalue={summary.ks_pvalue:.6g}"
+    )
+    print(f"summary: {summary_path}")
+    for path in plot_paths:
+        print(f"plot: {path}")
+
+
+if __name__ == "__main__":
+    main()