1 files changed, 96 insertions, 0 deletions
diff --git a/scripts/generate_paper_tables.py b/scripts/generate_paper_tables.py
index 95a67c8..3ef9ff7 100644
--- a/scripts/generate_paper_tables.py
+++ b/scripts/generate_paper_tables.py
@@ -194,6 +194,98 @@ def teacher_table(teacher_dir: Path, outdir: Path) -> None:
     write_fragment(outdir / "teacher_nested_se.tex", lines)
 
 
+TRAIN_SIZES = [64, 96, 128, 160, 192, 224, 256, 320]
+
+
+def _corr(xs: list[float], ys: list[float]) -> float:
+    mx, my = statistics.mean(xs), statistics.mean(ys)
+    cov = sum((x - mx) * (y - my) for x, y in zip(xs, ys))
+    return cov / math.sqrt(
+        sum((x - mx) ** 2 for x in xs) * sum((y - my) ** 2 for y in ys)
+    )
+
+
+def train_size_table(outputs: Path, outdir: Path) -> None:
+    lines = [
+        r"\begin{tabular}{@{}rcccc@{}}",
+        r"\toprule",
+        r"\(N\) & Measured & Frozen & Linear & Relin.\ \\",
+        r"\midrule",
+    ]
+    for n in TRAIN_SIZES:
+        rows = read_rows(
+            outputs
+            / f"compressed_operator_s20_256traj_T50_width64_N{n}"
+            / "compressed_operator_rows.csv"
+        )
+        measured = [float(r["empirical_gap"]) for r in rows]
+        maes = {
+            field: statistics.mean(
+                abs(float(r[field]) - m) for r, m in zip(rows, measured)
+            )
+            for field in ("fixed_gap", "linear_gap", "retangent_gap")
+        }
+        lines.append(
+            f"{n} & {statistics.mean(measured):.5f} & {maes['fixed_gap']:.5f} & "
+            f"{maes['linear_gap']:.5f} & {maes['retangent_gap']:.5f} \\\\"
+        )
+    lines += [r"\bottomrule", r"\end{tabular}"]
+    write_fragment(outdir / "train_size_predictors.tex", lines)
+
+
+def snapshot_location_table(outputs: Path, outdir: Path) -> None:
+    rows = [
+        r
+        for n in TRAIN_SIZES
+        for r in read_rows(
+            outputs
+            / f"compressed_operator_retangent_T50_width64_N{n}"
+            / "compressed_operator_rows.csv"
+        )
+    ]
+    lines = [
+        r"\begin{tabular}{@{}rccccc@{}}",
+        r"\toprule",
+        r"\(s\) & \(s/T\) & Frozen & Linear & Relin.\ & Corr.\ \\",
+        r"\midrule",
+    ]
+    for s in ("5", "10", "20"):
+        sub = [r for r in rows if r["early_steps"] == s]
+        measured = [float(r["empirical_gap"]) for r in sub]
+        maes = {
+            field: statistics.mean(
+                abs(float(r[field]) - m) for r, m in zip(sub, measured)
+            )
+            for field in ("fixed_gap", "linear_gap", "retangent_gap")
+        }
+        corr = _corr([float(r["linear_gap"]) for r in sub], measured)
+        lines.append(
+            f"{s} & {int(s)/50:.1f} & {maes['fixed_gap']:.5f} & "
+            f"{maes['linear_gap']:.5f} & {maes['retangent_gap']:.5f} & {corr:.5f} \\\\"
+        )
+    lines += [r"\bottomrule", r"\end{tabular}"]
+    write_fragment(outdir / "snapshot_location.tex", lines)
+
+
+def stress_grid_table(outputs: Path, outdir: Path) -> None:
+    rows = read_rows(outputs / "operator_stress_grid_summary" / "stress_grid_metrics.csv")
+    lines = [
+        r"\begin{tabular}{@{}rrrrrcccc@{}}",
+        r"\toprule",
+        r"\(d\) & \(w\) & \(T\) & \(s\) & Rows & Frozen MAE & Linear MAE & Relin.\ MAE & Linear corr.\ \\",
+        r"\midrule",
+    ]
+    for row in rows:
+        lines.append(
+            f"{int(row['hidden_layers'])} & {int(row['width'])} & "
+            f"{int(row['target_steps'])} & {int(row['early_steps'])} & {int(row['rows'])} & "
+            f"{float(row['fixed_mae']):.5f} & {float(row['linear_mae']):.5f} & "
+            f"{float(row['retangent_mae']):.5f} & {float(row['linear_corr']):.5f} \\\\"
+        )
+    lines += [r"\bottomrule", r"\end{tabular}"]
+    write_fragment(outdir / "stress_grid.tex", lines)
+
+
 def main() -> None:
     args = parse_args()
     initialization_table(args.depth_dir, args.outdir)
@@ -201,6 +293,10 @@ def main() -> None:
     cnn_table(args.cnn_dir, args.outdir)
     learning_rate_table(args.lr_dir, args.outdir)
     teacher_table(args.teacher_dir, args.outdir)
+    outputs = args.depth_dir.parent
+    train_size_table(outputs, args.outdir)
+    snapshot_location_table(outputs, args.outdir)
+    stress_grid_table(outputs, args.outdir)
 
 
 if __name__ == "__main__":