From 96e201556ac94057a4a5c8864cf422ad43d72d58 Mon Sep 17 00:00:00 2001
From: YurenHao0426 <Blackhao0426@gmail.com>
Date: Fri, 29 May 2026 08:48:14 -0500
Subject: Add initialization coverage distribution matching

---
 notes/02_experiment_notes.md                    |  58 +++
 notes/03_paper_outline.md                       |  13 +
 scripts/README.md                               |  30 ++
 scripts/initialization_distribution_matching.py | 457 ++++++++++++++++++++++++
 4 files changed, 558 insertions(+)
 create mode 100755 scripts/initialization_distribution_matching.py

diff --git a/notes/02_experiment_notes.md b/notes/02_experiment_notes.md
index 2fad18f..c7e7168 100644
--- a/notes/02_experiment_notes.md
+++ b/notes/02_experiment_notes.md
@@ -385,6 +385,64 @@ Random-target means remain close to \(1/D\) for all distributions, but worst-cas
 
 This empirically illustrates the prior-free minimax theorem: without target or weight prior information, anisotropic feedback cannot improve the worst-case angular bound.
 
+## Initialization Coverage Distribution Matching Run Log
+
+Script:
+
+```bash
+python scripts/initialization_distribution_matching.py --dimension 128 --target-samples 100000 --feedback-samples 100000 --batch-size 8192 --seed 2026 --subspace-dim 8 --anisotropy 64 --plot
+```
+
+Theory:
+
+For a feedback initialization distribution \(\mu\):
+
+\[
+M_\mu=\mathbb E[\hat b\hat b^\top].
+\]
+
+For a random target direction \(a\):
+
+\[
+A(a)=a^\top M_\mu a.
+\]
+
+If \(\lambda_i\) are eigenvalues of \(M_\mu\), then:
+
+\[
+A(a)
+=
+\frac{\sum_i \lambda_i G_i}{\sum_i G_i},
+\qquad
+G_i\sim\chi^2_1.
+\]
+
+Setup:
+
+- dimension \(D=128\)
+- target samples: `100000`
+- feedback samples for empirical \(M_\mu\): `100000`
+- schemes: `isotropic`, `rademacher`, `subspace`, `axis`, `geometric_axis`
+- subspace dimension: `8`
+- geometric axis anisotropy: `64`
+
+Results:
+
+| scheme | population \(\lambda_{\min}\) | empirical \(\lambda_{\min}\) | predicted mean | empirical mean | predicted std | empirical std | KS |
+|---|---:|---:|---:|---:|---:|---:|---:|
+| `isotropic` | `0.0078125` | `0.0072928` | `0.0078125` | `0.0078124` | `~0` | `0.0000346` | `nan` |
+| `rademacher` | `0.0078125` | `0.0072982` | `0.0078125` | `0.0078125` | `~0` | `0.0000342` | `nan` |
+| `subspace` | `0` | `0` | `0.0077967` | `0.0078031` | `0.0037405` | `0.0037335` | `0.00299` |
+| `axis` | `0` | `0` | `0.0078396` | `0.0078514` | `0.0109834` | `0.0109230` | `0.00351` |
+| `geometric_axis` | `0.0005111` | `0.0004100` | `0.0078116` | `0.0078155` | `0.0010441` | `0.0010425` | `0.00470` |
+
+Interpretation:
+
+- All schemes have mean coverage near \(1/D\), as expected from averaging over uniformly random targets.
+- Isotropic and rademacher have population point-mass coverage at \(1/D\); empirical spread comes from finite-sample \(M_\mu\) estimation noise, so KS against a point mass is not meaningful.
+- Non-isotropic schemes match the predicted target-coverage distribution tightly, with max nondegenerate KS `0.00470`.
+- The distribution shape, not just the mean, exposes the tradeoff: subspace and axis schemes keep the same average but create zero-coverage worst-case directions.
+
 ## Functional Capacity Overlap Run Log
 
 Script:
diff --git a/notes/03_paper_outline.md b/notes/03_paper_outline.md
index 328871f..6055722 100644
--- a/notes/03_paper_outline.md
+++ b/notes/03_paper_outline.md
@@ -128,6 +128,19 @@ Prior-aware corollary:
 \operatorname{tr}(\Sigma_A M_\mu).
 \]
 
+Distributional validation target:
+
+\[
+A(a)=a^\top M_\mu a
+=
+\frac{\sum_i \lambda_i G_i}{\sum_i G_i},
+\qquad
+G_i\sim\chi^2_1.
+\]
+
+Use this to compare predicted target-coverage distributions for isotropic,
+subspace, axis, and anisotropic feedback initializations.
+
 ## 7. Trajectory Bridge
 
 Introduce local approximation:
diff --git a/scripts/README.md b/scripts/README.md
index efd4208..02b55b0 100644
--- a/scripts/README.md
+++ b/scripts/README.md
@@ -143,6 +143,36 @@ The prior-free minimax theorem says:
 
 with equality for isotropic feedback. Outputs are written under `outputs/minimax_initialization/`.
 
+## Initialization Coverage Distribution Matching
+
+Run:
+
+```bash
+python scripts/initialization_distribution_matching.py --dimension 128 --target-samples 100000 --feedback-samples 100000 --batch-size 8192 --seed 2026 --subspace-dim 8 --anisotropy 64 --plot
+```
+
+For an initialization distribution \(\mu\), define:
+
+\[
+M_\mu=\mathbb E[\hat b\hat b^\top].
+\]
+
+For a random target direction \(a\), theory predicts the coverage distribution:
+
+\[
+A(a)=a^\top M_\mu a
+=
+\frac{\sum_i \lambda_i G_i}{\sum_i G_i},
+\qquad
+G_i\sim\chi^2_1,
+\]
+
+where \(\lambda_i\) are eigenvalues of \(M_\mu\). The script compares this
+population-predicted distribution with the distribution induced by an empirical
+second-moment estimate from sampled feedback directions.
+
+Outputs are written under `outputs/initialization_distribution_matching/`.
+
 ## Functional Capacity Overlap
 
 Run:
diff --git a/scripts/initialization_distribution_matching.py b/scripts/initialization_distribution_matching.py
new file mode 100755
index 0000000..f4a14d5
--- /dev/null
+++ b/scripts/initialization_distribution_matching.py
@@ -0,0 +1,457 @@
+#!/usr/bin/env python3
+"""Match feedback-initialization coverage distributions from M_mu spectra.
+
+For a feedback initialization distribution mu, define
+
+    M_mu = E[b_hat b_hat^T].
+
+For a uniformly random target direction a, the expected squared alignment over
+feedback draws is A(a) = a^T M_mu a. If lambda_i are eigenvalues of M_mu, then
+
+    A(a) = sum_i lambda_i G_i / sum_i G_i,  G_i ~ chi^2_1.
+
+This script compares the population-predicted distribution from known spectra
+with the empirical distribution obtained by estimating M_mu from sampled
+feedback directions.
+"""
+
+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
+
+
+@dataclass(frozen=True)
+class RunConfig:
+    dimension: int
+    target_samples: int
+    feedback_samples: int
+    batch_size: int
+    seed: int
+    schemes: list[str]
+    subspace_dim: int
+    anisotropy: float
+    outdir: str
+    plot: bool
+
+
+@dataclass(frozen=True)
+class MatchRow:
+    scheme: str
+    dimension: int
+    target_samples: int
+    feedback_samples: int
+    population_lambda_min: float
+    empirical_lambda_min: float
+    population_lambda_max: float
+    empirical_lambda_max: float
+    predicted_mean: float
+    empirical_mean: float
+    predicted_std: float
+    empirical_std: float
+    predicted_q01: float
+    empirical_q01: float
+    predicted_q50: float
+    empirical_q50: float
+    predicted_q99: float
+    empirical_q99: float
+    ks_2sample_statistic: float
+    ks_2sample_pvalue: float
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description="Distribution matching for feedback initialization coverage."
+    )
+    parser.add_argument("--dimension", type=int, default=128)
+    parser.add_argument("--target-samples", type=int, default=100_000)
+    parser.add_argument("--feedback-samples", type=int, default=100_000)
+    parser.add_argument("--batch-size", type=int, default=8192)
+    parser.add_argument("--seed", type=int, default=0)
+    parser.add_argument(
+        "--schemes",
+        nargs="+",
+        default=["isotropic", "rademacher", "subspace", "axis", "geometric_axis"],
+        choices=["isotropic", "rademacher", "subspace", "axis", "geometric_axis"],
+    )
+    parser.add_argument("--subspace-dim", type=int, default=8)
+    parser.add_argument("--anisotropy", type=float, default=64.0)
+    parser.add_argument(
+        "--outdir",
+        type=Path,
+        default=Path("outputs/initialization_distribution_matching"),
+    )
+    parser.add_argument("--plot", action="store_true")
+    return parser.parse_args()
+
+
+def validate_config(config: RunConfig) -> None:
+    if config.dimension < 2:
+        raise ValueError("--dimension must be at least 2.")
+    if config.target_samples < 1:
+        raise ValueError("--target-samples must be positive.")
+    if config.feedback_samples < 1:
+        raise ValueError("--feedback-samples must be positive.")
+    if config.batch_size < 1:
+        raise ValueError("--batch-size must be positive.")
+    if config.subspace_dim < 1:
+        raise ValueError("--subspace-dim must be positive.")
+    if config.anisotropy < 1:
+        raise ValueError("--anisotropy must be at least 1.")
+
+
+def population_eigenvalues(
+    scheme: str, dimension: int, subspace_dim: int, anisotropy: float
+) -> np.ndarray:
+    if scheme in {"isotropic", "rademacher"}:
+        return np.full(dimension, 1.0 / dimension)
+
+    if scheme == "axis":
+        values = np.zeros(dimension, dtype=np.float64)
+        values[-1] = 1.0
+        return values
+
+    if scheme == "subspace":
+        active = min(subspace_dim, dimension)
+        values = np.zeros(dimension, dtype=np.float64)
+        values[-active:] = 1.0 / active
+        return values
+
+    if scheme == "geometric_axis":
+        weights = np.geomspace(1.0, anisotropy, num=dimension)
+        return weights / weights.sum()
+
+    raise ValueError(f"Unknown scheme: {scheme}")
+
+
+def sample_feedback(
+    rng: np.random.Generator,
+    scheme: str,
+    count: int,
+    dimension: int,
+    eigenvalues: np.ndarray,
+    subspace_dim: int,
+) -> np.ndarray:
+    if scheme == "isotropic":
+        values = rng.standard_normal((count, dimension))
+        norms = np.linalg.norm(values, axis=1, keepdims=True)
+        return values / np.maximum(norms, np.finfo(float).tiny)
+
+    if scheme == "rademacher":
+        values = rng.choice(np.array([-1.0, 1.0]), size=(count, dimension))
+        return values / np.sqrt(dimension)
+
+    if scheme == "axis":
+        values = np.zeros((count, dimension), dtype=np.float64)
+        values[:, -1] = rng.choice(np.array([-1.0, 1.0]), size=count)
+        return values
+
+    if scheme == "subspace":
+        active = min(subspace_dim, dimension)
+        values = np.zeros((count, dimension), dtype=np.float64)
+        block = rng.standard_normal((count, active))
+        norms = np.linalg.norm(block, axis=1, keepdims=True)
+        values[:, -active:] = block / np.maximum(norms, np.finfo(float).tiny)
+        return values
+
+    if scheme == "geometric_axis":
+        indices = rng.choice(dimension, size=count, p=eigenvalues)
+        values = np.zeros((count, dimension), dtype=np.float64)
+        signs = rng.choice(np.array([-1.0, 1.0]), size=count)
+        values[np.arange(count), indices] = signs
+        return values
+
+    raise ValueError(f"Unknown scheme: {scheme}")
+
+
+def estimate_moment(
+    rng: np.random.Generator,
+    scheme: str,
+    dimension: int,
+    eigenvalues: np.ndarray,
+    subspace_dim: int,
+    feedback_samples: int,
+    batch_size: int,
+) -> np.ndarray:
+    moment = np.zeros((dimension, dimension), dtype=np.float64)
+    offset = 0
+    while offset < feedback_samples:
+        count = min(batch_size, feedback_samples - offset)
+        feedback = sample_feedback(
+            rng, scheme, count, dimension, eigenvalues, subspace_dim
+        )
+        moment += feedback.T @ feedback
+        offset += count
+    return moment / feedback_samples
+
+
+def target_coverage_from_eigenvalues(
+    rng: np.random.Generator, eigenvalues: np.ndarray, samples: int, batch_size: int
+) -> np.ndarray:
+    values = np.empty(samples, dtype=np.float64)
+    dimension = len(eigenvalues)
+    offset = 0
+    while offset < samples:
+        count = min(batch_size, samples - offset)
+        g = rng.chisquare(df=1.0, size=(count, dimension))
+        values[offset : offset + count] = (g @ eigenvalues) / np.sum(g, axis=1)
+        offset += count
+    return values
+
+
+def summarize(
+    scheme: str,
+    dimension: int,
+    target_samples: int,
+    feedback_samples: int,
+    population_eigs: np.ndarray,
+    empirical_eigs: np.ndarray,
+    predicted: np.ndarray,
+    empirical: np.ndarray,
+) -> MatchRow:
+    if np.std(predicted) < 1e-14:
+        ks_statistic = float("nan")
+        ks_pvalue = float("nan")
+    else:
+        ks = stats.ks_2samp(predicted, empirical)
+        ks_statistic = float(ks.statistic)
+        ks_pvalue = float(ks.pvalue)
+    return MatchRow(
+        scheme=scheme,
+        dimension=dimension,
+        target_samples=target_samples,
+        feedback_samples=feedback_samples,
+        population_lambda_min=float(np.min(population_eigs)),
+        empirical_lambda_min=float(np.min(empirical_eigs)),
+        population_lambda_max=float(np.max(population_eigs)),
+        empirical_lambda_max=float(np.max(empirical_eigs)),
+        predicted_mean=float(np.mean(predicted)),
+        empirical_mean=float(np.mean(empirical)),
+        predicted_std=float(np.std(predicted, ddof=1)),
+        empirical_std=float(np.std(empirical, ddof=1)),
+        predicted_q01=float(np.quantile(predicted, 0.01)),
+        empirical_q01=float(np.quantile(empirical, 0.01)),
+        predicted_q50=float(np.quantile(predicted, 0.50)),
+        empirical_q50=float(np.quantile(empirical, 0.50)),
+        predicted_q99=float(np.quantile(predicted, 0.99)),
+        empirical_q99=float(np.quantile(empirical, 0.99)),
+        ks_2sample_statistic=ks_statistic,
+        ks_2sample_pvalue=ks_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[MatchRow],
+    spectra: dict[str, dict[str, list[float]]],
+    outdir: Path,
+) -> None:
+    outdir.mkdir(parents=True, exist_ok=True)
+    write_csv(outdir / "distribution_match.csv", rows)
+    payload = {
+        "config": asdict(config),
+        "distribution_match": [asdict(row) for row in rows],
+        "spectra": spectra,
+    }
+    (outdir / "summary.json").write_text(
+        json.dumps(payload, indent=2, sort_keys=True) + "\n"
+    )
+
+
+def save_plots(
+    rows: list[MatchRow],
+    distributions: dict[str, tuple[np.ndarray, np.ndarray]],
+    spectra: dict[str, dict[str, list[float]]],
+    outdir: Path,
+) -> list[Path]:
+    outdir.mkdir(parents=True, exist_ok=True)
+    paths: list[Path] = []
+
+    spectrum_path = outdir / "population_vs_empirical_spectra.png"
+    plt.figure(figsize=(7, 4.5))
+    for scheme, values in sorted(spectra.items()):
+        empirical = np.array(values["empirical"])
+        plt.plot(np.sort(empirical), label=scheme)
+    plt.xlabel("eigenvalue index")
+    plt.ylabel("empirical eigenvalue of M_mu")
+    plt.title("Estimated feedback second-moment spectra")
+    plt.legend()
+    plt.tight_layout()
+    plt.savefig(spectrum_path, dpi=180)
+    plt.close()
+    paths.append(spectrum_path)
+
+    hist_path = outdir / "coverage_histograms.png"
+    plt.figure(figsize=(8, 5))
+    for scheme, (predicted, empirical) in sorted(distributions.items()):
+        plt.hist(
+            empirical,
+            bins=80,
+            density=True,
+            histtype="step",
+            linewidth=1.5,
+            label=f"{scheme} empirical",
+        )
+    plt.xlabel("target coverage a^T M_mu a")
+    plt.ylabel("density")
+    plt.title("Coverage distributions from estimated M_mu")
+    plt.legend(fontsize=8)
+    plt.tight_layout()
+    plt.savefig(hist_path, dpi=180)
+    plt.close()
+    paths.append(hist_path)
+
+    for scheme, (predicted, empirical) in sorted(distributions.items()):
+        if np.std(predicted) < 1e-14:
+            continue
+        qq_path = outdir / f"qq_{scheme}.png"
+        probs = (np.arange(1, len(predicted) + 1) - 0.5) / len(predicted)
+        pred_q = np.quantile(predicted, probs)
+        emp_q = np.quantile(empirical, probs)
+        max_value = float(max(pred_q[-1], emp_q[-1]))
+        min_value = float(min(pred_q[0], emp_q[0]))
+        plt.figure(figsize=(5, 5))
+        plt.scatter(pred_q, emp_q, s=4, alpha=0.25)
+        plt.plot([min_value, max_value], [min_value, max_value], color="black", linewidth=1)
+        plt.xlabel("population-predicted quantile")
+        plt.ylabel("empirical-estimated quantile")
+        plt.title(f"Q-Q coverage: {scheme}")
+        plt.tight_layout()
+        plt.savefig(qq_path, dpi=180)
+        plt.close()
+        paths.append(qq_path)
+
+    worst_path = outdir / "worst_best_coverage.png"
+    names = [row.scheme for row in rows]
+    x = np.arange(len(names))
+    plt.figure(figsize=(7, 4.5))
+    plt.bar(
+        x - 0.18,
+        [row.empirical_lambda_min for row in rows],
+        width=0.36,
+        label="worst target",
+    )
+    plt.bar(
+        x + 0.18,
+        [row.empirical_lambda_max for row in rows],
+        width=0.36,
+        label="best target",
+    )
+    plt.axhline(1.0 / rows[0].dimension, color="black", linestyle="--", linewidth=1, label="1/D")
+    plt.xticks(x, names, rotation=25, ha="right")
+    plt.ylabel("expected squared alignment")
+    plt.title("Worst and best target coverage")
+    plt.legend()
+    plt.tight_layout()
+    plt.savefig(worst_path, dpi=180)
+    plt.close()
+    paths.append(worst_path)
+
+    return paths
+
+
+def parse_config(args: argparse.Namespace) -> RunConfig:
+    return RunConfig(
+        dimension=args.dimension,
+        target_samples=args.target_samples,
+        feedback_samples=args.feedback_samples,
+        batch_size=args.batch_size,
+        seed=args.seed,
+        schemes=args.schemes,
+        subspace_dim=args.subspace_dim,
+        anisotropy=args.anisotropy,
+        outdir=str(args.outdir),
+        plot=args.plot,
+    )
+
+
+def main() -> None:
+    args = parse_args()
+    config = parse_config(args)
+    validate_config(config)
+    rng = np.random.default_rng(config.seed)
+
+    rows: list[MatchRow] = []
+    spectra: dict[str, dict[str, list[float]]] = {}
+    distributions: dict[str, tuple[np.ndarray, np.ndarray]] = {}
+
+    for scheme in config.schemes:
+        population_eigs = population_eigenvalues(
+            scheme, config.dimension, config.subspace_dim, config.anisotropy
+        )
+        moment = estimate_moment(
+            rng,
+            scheme,
+            config.dimension,
+            population_eigs,
+            config.subspace_dim,
+            config.feedback_samples,
+            config.batch_size,
+        )
+        empirical_eigs = np.linalg.eigvalsh(moment)
+        predicted = target_coverage_from_eigenvalues(
+            rng, population_eigs, config.target_samples, config.batch_size
+        )
+        empirical = target_coverage_from_eigenvalues(
+            rng, empirical_eigs, config.target_samples, config.batch_size
+        )
+        row = summarize(
+            scheme,
+            config.dimension,
+            config.target_samples,
+            config.feedback_samples,
+            population_eigs,
+            empirical_eigs,
+            predicted,
+            empirical,
+        )
+        rows.append(row)
+        spectra[scheme] = {
+            "population": population_eigs.tolist(),
+            "empirical": empirical_eigs.tolist(),
+        }
+        distributions[scheme] = (predicted, empirical)
+        print(
+            f"{scheme}: "
+            f"lambda_min_pop={row.population_lambda_min:.6g}, "
+            f"lambda_min_emp={row.empirical_lambda_min:.6g}, "
+            f"KS={row.ks_2sample_statistic:.6g}, "
+            f"mean_pred={row.predicted_mean:.6g}, "
+            f"mean_emp={row.empirical_mean:.6g}"
+        )
+
+    outdir = Path(config.outdir)
+    write_outputs(config, rows, spectra, outdir)
+    plot_paths = save_plots(rows, distributions, spectra, outdir) if config.plot else []
+
+    finite_ks = [row.ks_2sample_statistic for row in rows if np.isfinite(row.ks_2sample_statistic)]
+    max_ks = max(finite_ks) if finite_ks else float("nan")
+    max_mean_error = max(abs(row.predicted_mean - row.empirical_mean) for row in rows)
+    print(f"max_ks_2sample: {max_ks:.8g}")
+    print(f"max_abs_mean_error: {max_mean_error:.8g}")
+    print(f"distribution_match: {outdir / 'distribution_match.csv'}")
+    for path in plot_paths:
+        print(f"plot: {path}")
+
+
+if __name__ == "__main__":
+    main()
-- 
cgit v1.2.3