Add capacity scaling calculations

1 files changed, 241 insertions, 0 deletions

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()