path: root/scripts/capacity_scaling.py

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