import argparse
import csv
import io
import math
import subprocess
import sys
from pathlib import Path
from typing import Optional
import matplotlib.pyplot as plt
def run_sweep(
repo: Path,
checkpoint: Optional[str],
lengths: list[int],
thresholds: list[float],
target_ms: float,
vocab_size: Optional[int],
embed_dim: Optional[int],
hidden_dim: Optional[int],
output_dim: Optional[int],
num_neighbors: Optional[int],
plasticity: Optional[bool],
) -> list[dict[str, str]]:
example = repo / "target" / "release" / "examples" / "fractal_tiebreak"
if not example.exists():
print("Building fractal_tiebreak example...", file=sys.stderr)
subprocess.run(
["cargo", "build", "--example", "fractal_tiebreak", "--release"],
cwd=repo,
check=True,
)
cmd = [
str(example),
"--lengths",
",".join(str(l) for l in lengths),
"--thresholds",
",".join(str(t) for t in thresholds),
"--target-ms",
str(target_ms),
]
if checkpoint:
cmd.extend(["--checkpoint", checkpoint])
for flag, value in [
("--vocab-size", vocab_size),
("--embed-dim", embed_dim),
("--hidden-dim", hidden_dim),
("--output-dim", output_dim),
("--num-neighbors", num_neighbors),
]:
if value is not None:
cmd.extend([flag, str(value)])
if plasticity is not None:
cmd.extend(["--plasticity", "true" if plasticity else "false"])
result = subprocess.run(cmd, cwd=repo, check=True, capture_output=True, text=True)
return list(csv.DictReader(io.StringIO(result.stdout)))
def plot(rows: list[dict[str, str]], lengths: list[int], output: Path) -> None:
by_length: dict[int, list[dict[str, str]]] = {l: [] for l in lengths}
for r in rows:
by_length[int(r["length"])].append(r)
n = len(lengths)
cols = math.ceil(math.sqrt(n))
rows_grid = math.ceil(n / cols)
fig, axes = plt.subplots(rows_grid, cols, figsize=(5 * cols, 4 * rows_grid), squeeze=False)
for idx, length in enumerate(lengths):
ax = axes[idx // cols][idx % cols]
subset = by_length[length]
thresholds = [float(r["threshold"]) for r in subset]
mse = [max(1e-16, float(r["mse"])) for r in subset]
hybrid_ns = [float(r["hybrid_ns"]) for r in subset]
tiebreak_ns = [float(r["tiebreak_ns"]) for r in subset]
speedup = [h / t for h, t in zip(hybrid_ns, tiebreak_ns)]
frac_geo = [float(r["fraction_geometric"]) for r in subset]
ax.scatter(mse, speedup, c=thresholds, cmap="coolwarm", edgecolors="black", linewidths=0.5, zorder=3)
ax.plot(mse, speedup, "-", color="gray", alpha=0.4, zorder=2)
ax.set_xscale("log")
ax.set_xlabel("logit MSE vs hybrid")
ax.set_ylabel("speedup over hybrid")
ax.set_title(f"L = {length}")
ax.grid(True, alpha=0.3, which="both")
for idx in range(n, cols * rows_grid):
axes[idx // cols][idx % cols].set_visible(False)
cbar_ax = fig.add_axes([0.92, 0.15, 0.02, 0.7])
sm = plt.cm.ScalarMappable(cmap="coolwarm", norm=plt.Normalize(0.0, 1.0))
sm.set_array([])
fig.colorbar(sm, cax=cbar_ax, label="max fractal dimension")
fig.suptitle("Fractal-Dimension Tiebreak — Speedup vs Divergence", y=1.02)
fig.tight_layout(rect=[0, 0, 0.9, 1])
fig.savefig(output, dpi=150, bbox_inches="tight")
print(f"Saved plot to {output}")
def main() -> None:
parser = argparse.ArgumentParser(description="Plot fractal tiebreak sweep")
parser.add_argument("--checkpoint", default=None, help="path to a flatten_params checkpoint")
parser.add_argument("--lengths", default="64,128,256,512", help="comma-separated context lengths")
parser.add_argument(
"--thresholds",
default="0.0,0.2,0.4,0.6,0.8,1.0,inf",
help="comma-separated max-fractal-dimension thresholds",
)
parser.add_argument("--target-ms", type=float, default=50.0, help="target benchmark time per mode")
parser.add_argument("--vocab-size", type=int, default=None, help="model vocabulary size")
parser.add_argument("--embed-dim", type=int, default=None, help="attention embedding dimension")
parser.add_argument("--hidden-dim", type=int, default=None, help="MLP hidden dimension")
parser.add_argument("--output-dim", type=int, default=None, help="MLP output dimension")
parser.add_argument("--num-neighbors", type=int, default=None, help="number of geometric neighbours")
parser.add_argument("--plasticity", action="store_true", help="enable learnable edge weights")
parser.add_argument("--csv", default=None, help="use an existing CSV instead of re-running")
parser.add_argument("-o", "--output", default="target/fractal_tiebreak.png", help="output PNG path")
args = parser.parse_args()
repo = Path(__file__).resolve().parent.parent
lengths = [int(x.strip()) for x in args.lengths.split(",")]
thresholds = [float(x.strip()) if x.strip().lower() not in ("inf", "infinity") else float("inf") for x in args.thresholds.split(",")]
if args.csv:
with open(args.csv, "r", newline="") as f:
rows = list(csv.DictReader(f))
else:
rows = run_sweep(
repo,
args.checkpoint,
lengths,
thresholds,
args.target_ms,
args.vocab_size,
args.embed_dim,
args.hidden_dim,
args.output_dim,
args.num_neighbors,
args.plasticity if args.plasticity else None,
)
output = Path(args.output)
output.parent.mkdir(parents=True, exist_ok=True)
plot(rows, lengths, output)
if __name__ == "__main__":
main()