diff options
Diffstat (limited to 'experiments')
| -rw-r--r-- | experiments/analyze_coupled_ladder_scaling.py | 488 |
1 files changed, 488 insertions, 0 deletions
diff --git a/experiments/analyze_coupled_ladder_scaling.py b/experiments/analyze_coupled_ladder_scaling.py new file mode 100644 index 0000000..9f36863 --- /dev/null +++ b/experiments/analyze_coupled_ladder_scaling.py @@ -0,0 +1,488 @@ +#!/usr/bin/env python3 +"""Analyze and plot the paired digital CLLN scaling ladder.""" + +from __future__ import annotations + +import argparse +import csv +import json +from pathlib import Path + +import matplotlib as mpl +import matplotlib.pyplot as plt +import numpy as np + + +METHOD_ORDER = ( + "clean", + "matched_noise", + "raw", + "constant", + "sdil", +) +DISPLAY = { + "clean": "Clean", + "matched_noise": "Same-RMS noise", + "raw": "Raw imperfection", + "constant": "Static calibration", + "sdil": "SDIL", +} +STYLE = { + "clean": dict(color="#222222", marker="^", linestyle=":"), + "matched_noise": dict(color="#CC79A7", marker="D", linestyle="--"), + "raw": dict(color="#D55E00", marker="X", linestyle="-"), + "constant": dict(color="#666666", marker="s", linestyle="--"), + "sdil": dict(color="#0072B2", marker="o", linestyle="-"), +} + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument( + "--core", type=Path, + default=Path("results/coupled_ladder/p1_imperfection_pilot.json")) + parser.add_argument( + "--baselines", type=Path, + default=Path("results/coupled_ladder/p1_bias_baseline_pilot.json")) + parser.add_argument( + "--output-analysis", type=Path, + default=Path("results/coupled_ladder/p1_scaling_analysis.json")) + parser.add_argument( + "--output-csv", type=Path, + default=Path("results/coupled_ladder/p1_scaling_source.csv")) + parser.add_argument( + "--output-figure", type=Path, + default=Path("results/figs/figure_clln_scaling_pilot")) + parser.add_argument("--bootstrap-replicates", type=int, default=20000) + parser.add_argument("--bootstrap-seed", type=int, default=20260829) + return parser.parse_args() + + +def merge_reports(core: dict, baselines: dict) -> list[dict]: + records = {} + for report in (core, baselines): + for record in report["records"]: + key = ( + record["side"], + record["task_index"], + record["device_seed"], + ) + if key not in records: + records[key] = { + key: value for key, value in record.items() + if key != "methods" + } + records[key]["methods"] = {} + overlap = set(records[key]["methods"]) & set(record["methods"]) + if overlap: + raise ValueError(f"duplicate method records: {sorted(overlap)}") + records[key]["methods"].update(record["methods"]) + merged = list(records.values()) + merged.sort(key=lambda record: ( + record["side"], record["task_index"], record["device_seed"])) + return merged + + +def task_cluster_values( + records: list[dict], side: int, method: str, metric: str +) -> tuple[np.ndarray, np.ndarray]: + task_indices = sorted({ + record["task_index"] for record in records + if record["side"] == side and method in record["methods"] + }) + values = [] + for task_index in task_indices: + device_values = [ + record["methods"][method][metric] + for record in records + if record["side"] == side + and record["task_index"] == task_index + and method in record["methods"] + ] + values.append(float(np.mean(device_values))) + return np.asarray(task_indices), np.asarray(values) + + +def bootstrap_mean( + values: np.ndarray, + rng: np.random.Generator, + replicates: int, +) -> tuple[float, list[float]]: + indices = rng.integers(0, len(values), size=(replicates, len(values))) + means = np.mean(values[indices], axis=1) + return float(np.mean(values)), [ + float(value) for value in np.percentile(means, (2.5, 97.5)) + ] + + +def task_size_matrix( + records: list[dict], sizes: list[int], method: str, metric: str +) -> tuple[np.ndarray, np.ndarray]: + common_tasks = None + by_size = {} + for side in sizes: + tasks, values = task_cluster_values(records, side, method, metric) + by_size[side] = dict(zip(tasks.tolist(), values.tolist())) + task_set = set(tasks.tolist()) + common_tasks = task_set if common_tasks is None else ( + common_tasks & task_set) + task_indices = np.asarray(sorted(common_tasks), dtype=int) + matrix = np.asarray([ + [by_size[side][int(task)] for side in sizes] + for task in task_indices + ]) + return task_indices, matrix + + +def slope_bootstrap( + records: list[dict], + sizes: list[int], + edges: np.ndarray, + *, + replicates: int, + seed: int, +) -> dict: + _, clean = task_size_matrix( + records, sizes, "clean", "classification_error") + _, constant = task_size_matrix( + records, sizes, "constant", "classification_error") + _, sdil = task_size_matrix( + records, sizes, "sdil", "classification_error") + x = np.log10(edges) + + def slope(matrix: np.ndarray) -> float: + return float(np.polyfit(x, np.mean(matrix, axis=0), 1)[0]) + + constant_excess = constant - clean + sdil_excess = sdil - clean + point_constant = slope(constant_excess) + point_sdil = slope(sdil_excess) + point_difference = point_constant - point_sdil + point_reduction = ( + 100.0 * point_difference / point_constant + if point_constant > 0.0 else None) + + rng = np.random.default_rng(seed) + constant_slopes = [] + sdil_slopes = [] + slope_differences = [] + relative_reductions = [] + for _ in range(replicates): + sample = rng.integers(0, len(clean), size=len(clean)) + constant_slope = slope(constant_excess[sample]) + sdil_slope = slope(sdil_excess[sample]) + constant_slopes.append(constant_slope) + sdil_slopes.append(sdil_slope) + slope_differences.append(constant_slope - sdil_slope) + if constant_slope > 1e-12: + relative_reductions.append( + 100.0 * (constant_slope - sdil_slope) / constant_slope) + + def interval(values: list[float]) -> list[float]: + return [ + float(value) + for value in np.percentile(values, (2.5, 97.5)) + ] + + return { + "x_axis": "log10 learnable edges", + "y_axis": "classification error excess over clean", + "static_calibration_slope": point_constant, + "static_calibration_slope_95ci": interval(constant_slopes), + "sdil_slope": point_sdil, + "sdil_slope_95ci": interval(sdil_slopes), + "paired_slope_difference": point_difference, + "paired_slope_difference_95ci": interval(slope_differences), + "relative_slope_reduction_percent": point_reduction, + "relative_slope_reduction_95ci": ( + interval(relative_reductions) if relative_reductions else None), + "relative_reduction_valid_bootstrap_fraction": float( + len(relative_reductions) / replicates), + } + + +def trace_summary( + records: list[dict], side: int, method: str +) -> list[dict]: + method_records = [ + record["methods"][method] + for record in records + if record["side"] == side and method in record["methods"] + ] + epochs = [record["epoch"] for record in method_records[0]["trace"]] + output = [] + for index, epoch in enumerate(epochs): + values = np.asarray([ + record["trace"][index]["classification_error"] + for record in method_records + ]) + output.append({ + "epoch": epoch, + "mean_classification_error": float(np.mean(values)), + "standard_error": float( + np.std(values, ddof=1) / np.sqrt(len(values)) + if len(values) > 1 else 0.0), + }) + return output + + +def build_analysis( + records: list[dict], *, replicates: int, seed: int +) -> dict: + sizes = sorted({record["side"] for record in records}) + edges = np.asarray([ + next( + record["learnable_edges"] for record in records + if record["side"] == side + ) for side in sizes + ]) + rng = np.random.default_rng(seed) + summaries = {} + for side in sizes: + side_summary = {} + for method in METHOD_ORDER: + metrics = {} + for metric in ( + "classification_error", + "classification_error_auc", + "reached_stable_zero_error", + ): + _, values = task_cluster_values( + records, side, method, metric) + mean, interval = bootstrap_mean( + values, rng, replicates) + metrics[metric] = { + "mean": mean, + "task_bootstrap_95ci": interval, + } + side_summary[method] = metrics + clean = side_summary["clean"]["classification_error"]["mean"] + raw = side_summary["raw"]["classification_error"]["mean"] + sdil = side_summary["sdil"]["classification_error"]["mean"] + side_summary["sdil_raw_to_clean_gap_closed"] = float( + (raw - sdil) / (raw - clean)) + summaries[str(side)] = { + "learnable_edges": int(edges[sizes.index(side)]), + "methods": side_summary, + } + return { + "analysis": "digital_coupled_ladder_scaling_pilot_analysis", + "confirmatory": False, + "bootstrap": { + "unit": "task; component draws averaged within task", + "task_clusters": len({ + record["task_index"] for record in records}), + "replicates": replicates, + "seed": seed, + "interval": "percentile 95%", + }, + "sizes": sizes, + "summaries": summaries, + "excess_error_scaling": slope_bootstrap( + records, + sizes, + edges, + replicates=replicates, + seed=seed + 1, + ), + "largest_grid_learning_curves": { + method: trace_summary(records, sizes[-1], method) + for method in METHOD_ORDER + }, + } + + +def write_csv(path: Path, analysis: dict) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("w", newline="") as stream: + writer = csv.writer(stream) + writer.writerow(( + "side", "learnable_edges", "method", "metric", "mean", + "ci_low", "ci_high", + )) + for side, side_summary in analysis["summaries"].items(): + for method in METHOD_ORDER: + for metric, values in side_summary["methods"][method].items(): + writer.writerow(( + side, + side_summary["learnable_edges"], + method, + metric, + values["mean"], + values["task_bootstrap_95ci"][0], + values["task_bootstrap_95ci"][1], + )) + + +def plot_figure(path: Path, analysis: dict) -> None: + mpl.rcParams.update({ + "font.family": "DejaVu Sans", + "font.size": 8.5, + "axes.labelsize": 9, + "axes.titlesize": 9.5, + "legend.fontsize": 7.5, + "xtick.labelsize": 8, + "ytick.labelsize": 8, + "axes.spines.top": False, + "axes.spines.right": False, + "svg.fonttype": "none", + "pdf.fonttype": 42, + "figure.facecolor": "white", + "axes.facecolor": "white", + }) + sizes = analysis["sizes"] + edges = np.asarray([ + analysis["summaries"][str(side)]["learnable_edges"] + for side in sizes + ]) + figure, axes = plt.subplots(1, 3, figsize=(10.6, 3.05)) + edge_labels = [ + f"{value / 1000.0:.1f}k" if value >= 1000 else str(value) + for value in edges + ] + + for method in METHOD_ORDER: + means = np.asarray([ + analysis["summaries"][str(side)]["methods"][method][ + "classification_error" + ]["mean"] * 100.0 + for side in sizes + ]) + intervals = np.asarray([ + analysis["summaries"][str(side)]["methods"][method][ + "classification_error" + ]["task_bootstrap_95ci"] + for side in sizes + ]) * 100.0 + axes[0].errorbar( + edges, + means, + yerr=np.vstack((means - intervals[:, 0], intervals[:, 1] - means)), + linewidth=1.7 if method == "sdil" else 1.15, + markersize=4.5, + capsize=2.0, + label=DISPLAY[method], + **STYLE[method], + ) + axes[0].set_xscale("log", base=2) + axes[0].set_xticks(edges, edge_labels) + axes[0].set_ylim(-3, 70) + axes[0].set_xlabel("Learnable edges") + axes[0].set_ylabel("Final classification error (%)") + reduction = analysis["excess_error_scaling"][ + "relative_slope_reduction_percent"] + axes[0].set_title( + f"(a) Error-growth slope reduced {reduction:.0f}% vs static calibration") + + for method in METHOD_ORDER: + means = np.asarray([ + analysis["summaries"][str(side)]["methods"][method][ + "reached_stable_zero_error" + ]["mean"] * 100.0 + for side in sizes + ]) + axes[1].plot( + edges, + means, + linewidth=1.7 if method == "sdil" else 1.15, + markersize=4.5, + label=DISPLAY[method], + **STYLE[method], + ) + axes[1].set_xscale("log", base=2) + axes[1].set_xticks(edges, edge_labels) + axes[1].set_ylim(-5, 105) + axes[1].set_xlabel("Learnable edges") + axes[1].set_ylabel("Stable zero-error runs (%)") + axes[1].set_title("(b) Recovery remains reliable") + + for method in METHOD_ORDER: + trace = analysis["largest_grid_learning_curves"][method] + epochs = np.asarray([record["epoch"] for record in trace]) + means = np.asarray([ + record["mean_classification_error"] * 100.0 + for record in trace + ]) + standard_errors = np.asarray([ + record["standard_error"] * 100.0 for record in trace + ]) + axes[2].plot( + epochs, + means, + linewidth=1.7 if method == "sdil" else 1.15, + label=DISPLAY[method], + **{key: value for key, value in STYLE[method].items() + if key != "marker"}, + ) + axes[2].fill_between( + epochs, + np.maximum(0.0, means - standard_errors), + means + standard_errors, + color=STYLE[method]["color"], + alpha=0.08, + linewidth=0, + ) + axes[2].set_xlim(0, 600) + axes[2].set_ylim(-3, 70) + axes[2].set_xlabel("Training epoch") + axes[2].set_ylabel("Classification error (%)") + axes[2].set_title("(c) Learning at 2,048 edges") + + for axis in axes: + axis.grid(axis="y", color="#D9D9D9", linewidth=0.55, alpha=0.8) + axis.tick_params(length=3) + handles, labels = axes[0].get_legend_handles_labels() + figure.legend( + handles, + labels, + loc="upper center", + ncol=5, + frameon=False, + bbox_to_anchor=(0.5, 1.02), + ) + figure.text( + 0.995, + 0.005, + "Exploratory pilot: 5 tasks, 1 component draw per task and size", + ha="right", + va="bottom", + fontsize=7, + color="#666666", + ) + figure.tight_layout(rect=(0.0, 0.04, 1.0, 0.92), w_pad=2.0) + path.parent.mkdir(parents=True, exist_ok=True) + figure.savefig(path.with_suffix(".svg"), bbox_inches="tight") + figure.savefig(path.with_suffix(".pdf"), bbox_inches="tight") + figure.savefig(path.with_suffix(".png"), dpi=240, bbox_inches="tight") + plt.close(figure) + + +def main() -> None: + args = parse_args() + core = json.loads(args.core.read_text()) + baselines = json.loads(args.baselines.read_text()) + records = merge_reports(core, baselines) + analysis = build_analysis( + records, + replicates=args.bootstrap_replicates, + seed=args.bootstrap_seed, + ) + analysis["sources"] = [str(args.core), str(args.baselines)] + args.output_analysis.parent.mkdir(parents=True, exist_ok=True) + args.output_analysis.write_text(json.dumps(analysis, indent=2) + "\n") + write_csv(args.output_csv, analysis) + plot_figure(args.output_figure, analysis) + print(json.dumps({ + "excess_error_scaling": analysis["excess_error_scaling"], + "largest_side": analysis["sizes"][-1], + "largest_side_methods": analysis["summaries"][ + str(analysis["sizes"][-1]) + ]["methods"], + }, indent=2)) + print(f"wrote {args.output_analysis}") + print(f"wrote {args.output_csv}") + print(f"wrote {args.output_figure}.svg/.pdf/.png") + + +if __name__ == "__main__": + main() |