Add capacity scaling calculations

3 files changed, 293 insertions, 0 deletions

diff --git a/notes/02_experiment_notes.md b/notes/02_experiment_notes.md
index b7f17a5..ab69cd8 100644
--- a/notes/02_experiment_notes.md
+++ b/notes/02_experiment_notes.md
@@ -235,3 +235,27 @@ Result:
 - KS p-value: `0.633285`
 
 This is a clean first-pass validation for the isotropic Gaussian case.
+
+## Scaling Run Log
+
+Script:
+
+```bash
+python scripts/capacity_scaling.py --plot
+```
+
+Default sweep:
+
+- widths: `16, 32, 64, 128`
+- feedback-aligned layer counts: `1, 2, 4, 8, 16`
+- fixed threshold: \(q=0.01\)
+- chance-level threshold: \(q=1/D\)
+- log unit: nats
+
+Result:
+
+- rows written: `40`
+- fixed-threshold max total cost: `1361.74` nats at width `128`, layers `16`
+- chance-level max total cost: `18.3652` nats at width `128`, layers `16`
+
+This cleanly separates the fixed-threshold regime, where total cost scales like \(Ln^2\), from the chance-level regime, where per-layer cost is nearly width-independent.
diff --git a/scripts/README.md b/scripts/README.md
index 1fe2ad9..3124261 100644
--- a/scripts/README.md
+++ b/scripts/README.md
@@ -25,3 +25,31 @@ D=\texttt{rows}\times\texttt{cols}.
 
 Outputs are written under `outputs/static_alignment_beta/`, which is ignored by Git.
 
+## Capacity Scaling
+
+Run:
+
+```bash
+python scripts/capacity_scaling.py --plot
+```
+
+This computes:
+
+\[
+C_l(q)
+=
+-\log\Pr(Q_l\ge q)
+\]
+
+for equal-width feedback blocks with \(D=n^2\), then sums over the number of feedback-aligned layers:
+
+\[
+C_{\mathrm{all}}=\sum_l C_l(q_l).
+\]
+
+The default run compares two regimes:
+
+- `fixed`: \(q=0.01\), where \(C_{\mathrm{all}}\) grows like \(\Theta(Ln^2)\).
+- `chance`: \(q=1/D\), where \(C_{\mathrm{all}}\) grows mostly with \(L\).
+
+Outputs are written under `outputs/capacity_scaling/`.
diff --git a/scripts/capacity_scaling.py b/scripts/capacity_scaling.py
new file mode 100755
index 0000000..2a4c220
--- /dev/null
+++ b/scripts/capacity_scaling.py
@@ -0,0 +1,241 @@
+#!/usr/bin/env python3
+"""Compute log-volume capacity scaling for FA alignment thresholds."""
+
+from __future__ import annotations
+
+import argparse
+import csv
+import json
+import math
+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 ScalingConfig:
+    widths: list[int]
+    feedback_layers: list[int]
+    threshold_modes: list[str]
+    fixed_q: float
+    chance_c: float
+    log_base: str
+    outdir: str
+    plot: bool
+
+
+@dataclass(frozen=True)
+class ScalingRow:
+    threshold_mode: str
+    width: int
+    feedback_layers: int
+    dimension: int
+    threshold_q: float
+    per_layer_cost: float
+    total_cost: float
+    x_ln2: int
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description="Compute FA log-volume capacity scaling tables."
+    )
+    parser.add_argument(
+        "--widths",
+        type=int,
+        nargs="+",
+        default=[16, 32, 64, 128],
+        help="Equal hidden widths to evaluate.",
+    )
+    parser.add_argument(
+        "--feedback-layers",
+        type=int,
+        nargs="+",
+        default=[1, 2, 4, 8, 16],
+        help="Number of equal-width feedback-aligned layer blocks.",
+    )
+    parser.add_argument(
+        "--threshold-mode",
+        choices=["fixed", "chance", "both"],
+        default="both",
+        help="fixed uses q; chance uses q=c/D.",
+    )
+    parser.add_argument(
+        "--q",
+        type=float,
+        default=0.01,
+        help="Fixed squared-alignment threshold.",
+    )
+    parser.add_argument(
+        "--c",
+        type=float,
+        default=1.0,
+        help="Chance-level multiplier for q=c/D.",
+    )
+    parser.add_argument(
+        "--log-base",
+        choices=["e", "2"],
+        default="e",
+        help="Use nats (e) or bits (2).",
+    )
+    parser.add_argument(
+        "--outdir",
+        type=Path,
+        default=Path("outputs/capacity_scaling"),
+        help="Directory for tables and plots.",
+    )
+    parser.add_argument("--plot", action="store_true", help="Save scaling plots.")
+    return parser.parse_args()
+
+
+def threshold_modes(mode: str) -> list[str]:
+    if mode == "both":
+        return ["fixed", "chance"]
+    return [mode]
+
+
+def log_base_value(log_base: str) -> float:
+    if log_base == "e":
+        return 1.0
+    if log_base == "2":
+        return math.log(2.0)
+    raise ValueError(f"Unknown log base: {log_base}")
+
+
+def per_layer_capacity_cost(dimension: int, q: float, log_base: str) -> float:
+    if not 0 <= q <= 1:
+        raise ValueError(f"Threshold must be in [0, 1], got {q}.")
+    alpha = 0.5
+    beta_param = (dimension - 1) / 2
+    log_tail = stats.beta(alpha, beta_param).logsf(q)
+    return float(-log_tail / log_base_value(log_base))
+
+
+def make_rows(config: ScalingConfig) -> list[ScalingRow]:
+    rows: list[ScalingRow] = []
+    for mode in config.threshold_modes:
+        for width in config.widths:
+            dimension = width * width
+            for layers in config.feedback_layers:
+                if mode == "fixed":
+                    q = config.fixed_q
+                elif mode == "chance":
+                    q = min(config.chance_c / dimension, 1.0)
+                else:
+                    raise ValueError(f"Unknown threshold mode: {mode}")
+
+                per_layer = per_layer_capacity_cost(dimension, q, config.log_base)
+                rows.append(
+                    ScalingRow(
+                        threshold_mode=mode,
+                        width=width,
+                        feedback_layers=layers,
+                        dimension=dimension,
+                        threshold_q=q,
+                        per_layer_cost=per_layer,
+                        total_cost=layers * per_layer,
+                        x_ln2=layers * dimension,
+                    )
+                )
+    return rows
+
+
+def write_outputs(config: ScalingConfig, rows: list[ScalingRow], outdir: Path) -> None:
+    outdir.mkdir(parents=True, exist_ok=True)
+
+    payload = {
+        "config": asdict(config),
+        "rows": [asdict(row) for row in rows],
+    }
+    (outdir / "scaling.json").write_text(
+        json.dumps(payload, indent=2, sort_keys=True) + "\n"
+    )
+
+    with (outdir / "scaling.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[ScalingRow], outdir: Path, log_base: str) -> list[Path]:
+    outdir.mkdir(parents=True, exist_ok=True)
+    paths: list[Path] = []
+    ylabel = "capacity cost (nats)" if log_base == "e" else "capacity cost (bits)"
+
+    for mode in sorted({row.threshold_mode for row in rows}):
+        mode_rows = [row for row in rows if row.threshold_mode == mode]
+        path = outdir / f"{mode}_scaling.png"
+
+        plt.figure(figsize=(7, 4.5))
+        for width in sorted({row.width for row in mode_rows}):
+            width_rows = sorted(
+                [row for row in mode_rows if row.width == width],
+                key=lambda row: row.feedback_layers,
+            )
+            x = [row.x_ln2 if mode == "fixed" else row.feedback_layers for row in width_rows]
+            y = [row.total_cost for row in width_rows]
+            plt.plot(x, y, marker="o", label=f"width={width}")
+
+        if mode == "fixed":
+            plt.xlabel("feedback_layers x width^2")
+        else:
+            plt.xlabel("feedback_layers")
+        plt.ylabel(ylabel)
+        plt.title(f"{mode} threshold capacity scaling")
+        plt.legend()
+        plt.tight_layout()
+        plt.savefig(path, dpi=180)
+        plt.close()
+        paths.append(path)
+
+    return paths
+
+
+def main() -> None:
+    args = parse_args()
+    if any(width < 2 for width in args.widths):
+        raise ValueError("All widths must be at least 2.")
+    if any(layers < 1 for layers in args.feedback_layers):
+        raise ValueError("All feedback layer counts must be positive.")
+    if not 0 <= args.q <= 1:
+        raise ValueError("--q must be in [0, 1].")
+    if args.c < 0:
+        raise ValueError("--c must be non-negative.")
+
+    config = ScalingConfig(
+        widths=args.widths,
+        feedback_layers=args.feedback_layers,
+        threshold_modes=threshold_modes(args.threshold_mode),
+        fixed_q=args.q,
+        chance_c=args.c,
+        log_base=args.log_base,
+        outdir=str(args.outdir),
+        plot=args.plot,
+    )
+
+    rows = make_rows(config)
+    write_outputs(config, rows, args.outdir)
+    plot_paths = save_plots(rows, args.outdir, args.log_base) if args.plot else []
+
+    print(f"rows: {len(rows)}")
+    print(f"table: {args.outdir / 'scaling.csv'}")
+    print(f"json: {args.outdir / 'scaling.json'}")
+    for path in plot_paths:
+        print(f"plot: {path}")
+
+    for mode in config.threshold_modes:
+        subset = [row for row in rows if row.threshold_mode == mode]
+        max_row = max(subset, key=lambda row: row.total_cost)
+        print(
+            f"{mode}: max total_cost={max_row.total_cost:.6g} "
+            f"at width={max_row.width}, layers={max_row.feedback_layers}, "
+            f"q={max_row.threshold_q:.6g}"
+        )
+
+
+if __name__ == "__main__":
+    main()