#!/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()