diff options
Diffstat (limited to 'scripts/functional_capacity_overlap.py')
| -rwxr-xr-x | scripts/functional_capacity_overlap.py | 298 |
1 files changed, 298 insertions, 0 deletions
diff --git a/scripts/functional_capacity_overlap.py b/scripts/functional_capacity_overlap.py new file mode 100755 index 0000000..266f2ea --- /dev/null +++ b/scripts/functional_capacity_overlap.py @@ -0,0 +1,298 @@ +#!/usr/bin/env python3 +"""Validate random-subspace functional capacity overlap formulas. + +Let S be a d-dimensional task-sensitive subspace in R^P and E be a k-dimensional +alignment constraint subspace. For generic subspaces: + + hard rank loss = max(0, k - (P - d)) + +and the soft overlap T = tr(P_E P_S) satisfies: + + E[T] = kd/P. +""" + +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 + + +@dataclass(frozen=True) +class RunConfig: + parameters: int + task_rank: int + constraint_ranks: list[int] + trials: int + seed: int + rank_tol: float + outdir: str + plot: bool + + +@dataclass(frozen=True) +class CapacityRow: + constraint_rank: int + hard_loss_empirical_mean: float + hard_loss_empirical_std: float + hard_loss_theory: float + soft_overlap_empirical_mean: float + soft_overlap_empirical_var: float + soft_overlap_theory_mean: float + soft_overlap_theory_var: float + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Validate redundancy-exhaustion capacity formulas." + ) + parser.add_argument( + "--parameters", + type=int, + default=128, + help="Ambient parameter dimension P.", + ) + parser.add_argument( + "--task-rank", + type=int, + default=32, + help="Task-sensitive subspace dimension d.", + ) + parser.add_argument( + "--constraint-ranks", + type=int, + nargs="+", + default=[0, 16, 32, 64, 96, 112, 128], + help="Alignment constraint ranks k to test.", + ) + parser.add_argument("--trials", type=int, default=200, help="Monte Carlo trials.") + parser.add_argument("--seed", type=int, default=0, help="Random seed.") + parser.add_argument( + "--rank-tol", + type=float, + default=1e-9, + help="Numerical tolerance for rank estimation.", + ) + parser.add_argument( + "--outdir", + type=Path, + default=Path("outputs/functional_capacity_overlap"), + help="Output directory.", + ) + parser.add_argument("--plot", action="store_true", help="Save plots.") + return parser.parse_args() + + +def random_orthonormal( + rng: np.random.Generator, ambient_dim: int, subspace_dim: int +) -> np.ndarray: + if subspace_dim == 0: + return np.empty((ambient_dim, 0), dtype=np.float64) + values = rng.standard_normal((ambient_dim, subspace_dim)) + q, _ = np.linalg.qr(values, mode="reduced") + return q[:, :subspace_dim] + + +def hard_rank_loss( + task_basis: np.ndarray, constraint_basis: np.ndarray, rank_tol: float +) -> int: + task_rank = task_basis.shape[1] + if constraint_basis.shape[1] == 0: + return 0 + projected = task_basis - constraint_basis @ (constraint_basis.T @ task_basis) + remaining_rank = np.linalg.matrix_rank(projected, tol=rank_tol) + return int(task_rank - remaining_rank) + + +def soft_overlap(task_basis: np.ndarray, constraint_basis: np.ndarray) -> float: + if task_basis.shape[1] == 0 or constraint_basis.shape[1] == 0: + return 0.0 + cross = constraint_basis.T @ task_basis + return float(np.sum(cross * cross)) + + +def theory_hard_loss(parameters: int, task_rank: int, constraint_rank: int) -> int: + return max(0, constraint_rank - (parameters - task_rank)) + + +def theory_soft_variance(parameters: int, task_rank: int, constraint_rank: int) -> float: + p = parameters + d = task_rank + k = constraint_rank + if p <= 2: + return 0.0 + numerator = 2 * k * d * (p - k) * (p - d) + denominator = p * p * (p - 1) * (p + 2) + return numerator / denominator + + +def run_trials(config: RunConfig) -> list[CapacityRow]: + rng = np.random.default_rng(config.seed) + rows: list[CapacityRow] = [] + + for constraint_rank in config.constraint_ranks: + hard_losses = np.empty(config.trials, dtype=np.float64) + overlaps = np.empty(config.trials, dtype=np.float64) + + for trial in range(config.trials): + task_basis = random_orthonormal( + rng, config.parameters, config.task_rank + ) + constraint_basis = random_orthonormal( + rng, config.parameters, constraint_rank + ) + hard_losses[trial] = hard_rank_loss( + task_basis, constraint_basis, config.rank_tol + ) + overlaps[trial] = soft_overlap(task_basis, constraint_basis) + + rows.append( + CapacityRow( + constraint_rank=constraint_rank, + hard_loss_empirical_mean=float(np.mean(hard_losses)), + hard_loss_empirical_std=float(np.std(hard_losses, ddof=1)) + if config.trials > 1 + else 0.0, + hard_loss_theory=float( + theory_hard_loss( + config.parameters, config.task_rank, constraint_rank + ) + ), + soft_overlap_empirical_mean=float(np.mean(overlaps)), + soft_overlap_empirical_var=float(np.var(overlaps, ddof=1)) + if config.trials > 1 + else 0.0, + soft_overlap_theory_mean=constraint_rank + * config.task_rank + / config.parameters, + soft_overlap_theory_var=theory_soft_variance( + config.parameters, config.task_rank, constraint_rank + ), + ) + ) + + return rows + + +def write_outputs(config: RunConfig, rows: list[CapacityRow], outdir: Path) -> None: + outdir.mkdir(parents=True, exist_ok=True) + payload = { + "config": asdict(config), + "rows": [asdict(row) for row in rows], + } + (outdir / "summary.json").write_text( + json.dumps(payload, indent=2, sort_keys=True) + "\n" + ) + + with (outdir / "summary.csv").open("w", newline="") as handle: + writer = csv.DictWriter(handle, fieldnames=list(asdict(rows[0]).keys())) + writer.writeheader() + for row in rows: + writer.writerow(asdict(row)) + + +def save_plots(rows: list[CapacityRow], outdir: Path) -> list[Path]: + outdir.mkdir(parents=True, exist_ok=True) + paths: list[Path] = [] + k = np.array([row.constraint_rank for row in rows]) + + hard_path = outdir / "hard_rank_loss.png" + plt.figure(figsize=(7, 4.5)) + plt.plot( + k, + [row.hard_loss_empirical_mean for row in rows], + marker="o", + label="empirical", + ) + plt.plot( + k, + [row.hard_loss_theory for row in rows], + linestyle="--", + color="black", + label="theory", + ) + plt.xlabel("constraint rank k") + plt.ylabel("hard functional rank loss") + plt.title("Redundancy-exhaustion transition") + plt.legend() + plt.tight_layout() + plt.savefig(hard_path, dpi=180) + plt.close() + paths.append(hard_path) + + soft_path = outdir / "soft_overlap.png" + plt.figure(figsize=(7, 4.5)) + plt.plot( + k, + [row.soft_overlap_empirical_mean for row in rows], + marker="o", + label="empirical mean", + ) + plt.plot( + k, + [row.soft_overlap_theory_mean for row in rows], + linestyle="--", + color="black", + label="theory mean", + ) + plt.xlabel("constraint rank k") + plt.ylabel("soft overlap tr(P_E P_S)") + plt.title("Soft functional capacity overlap") + plt.legend() + plt.tight_layout() + plt.savefig(soft_path, dpi=180) + plt.close() + paths.append(soft_path) + + return paths + + +def main() -> None: + args = parse_args() + if args.parameters < 1: + raise ValueError("--parameters must be positive.") + if not 0 <= args.task_rank <= args.parameters: + raise ValueError("--task-rank must be in [0, parameters].") + if any(k < 0 or k > args.parameters for k in args.constraint_ranks): + raise ValueError("All constraint ranks must be in [0, parameters].") + if args.trials < 1: + raise ValueError("--trials must be positive.") + + config = RunConfig( + parameters=args.parameters, + task_rank=args.task_rank, + constraint_ranks=args.constraint_ranks, + trials=args.trials, + seed=args.seed, + rank_tol=args.rank_tol, + outdir=str(args.outdir), + plot=args.plot, + ) + + rows = run_trials(config) + write_outputs(config, rows, args.outdir) + plot_paths = save_plots(rows, args.outdir) if args.plot else [] + + print(f"parameters P: {args.parameters}") + print(f"task rank d: {args.task_rank}") + print(f"trials: {args.trials}") + print(f"table: {args.outdir / 'summary.csv'}") + for row in rows: + print( + f"k={row.constraint_rank}: " + f"hard={row.hard_loss_empirical_mean:.6g} " + f"(theory {row.hard_loss_theory:.6g}), " + f"soft={row.soft_overlap_empirical_mean:.6g} " + f"(theory {row.soft_overlap_theory_mean:.6g})" + ) + for path in plot_paths: + print(f"plot: {path}") + + +if __name__ == "__main__": + main() |