figure fixes: budget-map label placement, bar labels, schematic loop

make_report.py

@@ -160,6 +160,11 @@ def plot_training(evals, out):
     fig.tight_layout(); fig.savefig(out, dpi=150); plt.close(fig)
 
 
+SHORT = {"fixed4x1": "fixed-4", "rec1x4": "rec-1x4", "rec2x2": "rec-2x2",
+         "rec1xdynamic": "rec-1x-dyn", "rec4x4": "rec-4x4",
+         "fixed4-long": "fixed-4-long"}
+
+
 def plot_headline_bars(evals, out):
     """Bar-chart companion to the headline table. Two panels:
     (a) the matched-FLOP family, (b) the iso-param / iso-FLOP trio."""
@@ -188,7 +193,7 @@ def plot_headline_bars(evals, out):
         ptxt = []
         for f in fams:
             p, fl = budget_of(evals[f]) if evals[f] else (0, 0)
-            ptxt.append(f"{LABELS[f].split(' (')[0]}\n"
+            ptxt.append(f"{SHORT[f]}\n"
                         f"{p/1000:.0f}k params\n{fl:.2e} FLOPs")
         ax.set_xticks(x); ax.set_xticklabels(ptxt, fontsize=7)
         ax.axhline(1 / 64, color="gray", ls=":", lw=1)
@@ -205,25 +210,32 @@ def plot_headline_bars(evals, out):
 def plot_budget_map(evals, out):
     """Where each family sits in the (params, total training FLOPs) plane,
     with the two matching axes drawn in."""
-    fig, ax = plt.subplots(figsize=(7.5, 5))
+    fig, ax = plt.subplots(figsize=(7.5, 5.5))
+    # hand-placed label offsets (points) so nothing overlaps
+    OFF = {"fixed4x1": (8, -14, "left"), "rec1x4": (0, -16, "center"),
+           "rec2x2": (0, 10, "center"), "rec1xdynamic": (8, 6, "left"),
+           "rec4x4": (-8, 10, "right"), "fixed4-long": (8, -14, "left")}
     for f in FAMILIES:
         if not evals[f]:
             continue
         p, fl = budget_of(evals[f])
         ax.scatter(p, fl, s=90, zorder=3)
-        ax.annotate(LABELS[f], (p, fl), fontsize=7,
-                    xytext=(6, 4), textcoords="offset points")
+        dx, dy, ha = OFF[f]
+        ax.annotate(LABELS[f], (p, fl), fontsize=7, ha=ha,
+                    xytext=(dx, dy), textcoords="offset points")
     ax.axvline(829696, color="gray", ls="--", lw=1)
-    ax.annotate("iso-parameter pair\n(fixed-4 vs rec-4x4)", (829696, 1.5e15),
+    ax.annotate("iso-parameter pair\n(fixed-4 vs rec-4x4)", (829696, 1.45e15),
                 fontsize=7, color="gray", ha="right")
     ax.axhline(2.185e15, color="gray", ls=":", lw=1)
-    ax.annotate("matched-FLOP family\n(same per-step and total FLOPs)",
-                (2.4e5, 2.185e15), fontsize=7, color="gray",
-                xytext=(6, 5), textcoords="offset points")
+    ax.annotate("matched-FLOP family (same per-step and total FLOPs)",
+                (4.8e5, 2.185e15), fontsize=7, color="gray",
+                xytext=(0, 7), textcoords="offset points", ha="center")
     ax.axhline(8.72e15, color="gray", ls="-.", lw=1)
-    ax.annotate("iso-FLOP pair\n(rec-4x4 vs fixed-4-long)",
-                (2.4e5, 8.72e15), fontsize=7, color="gray",
-                xytext=(6, 5), textcoords="offset points")
+    ax.annotate("iso-FLOP pair (rec-4x4 vs fixed-4-long)",
+                (5.0e5, 8.72e15), fontsize=7, color="gray",
+                xytext=(0, -13), textcoords="offset points", ha="center")
+    ax.set_xlim(1.8e5, 1.6e6)
+    ax.set_ylim(1.2e15, 1.1e16)
     ax.set_xscale("log"); ax.set_yscale("log")
     ax.set_xlabel("parameters"); ax.set_ylabel("total training FLOPs")
     ax.set_title("Experimental design: what is matched against what")
@@ -267,11 +279,11 @@ def plot_schematic(out):
     box(ax, 1.2, 3.2, 1.6, 0.8, "shared blocks\n(same weights)", "tab:orange")
     box(ax, 1.2, 1.8, 1.6, 0.6, "readout head", "lightyellow")
     arrow(ax, 2.0, 3.2, 2.4)
-    ax.annotate("", xy=(1.0, 3.6), xytext=(1.0, 4.4),
+    ax.annotate("", xy=(2.8, 3.65), xytext=(2.8, 3.35),
                 arrowprops=dict(arrowstyle="-|>", connectionstyle=
-                                "arc3,rad=0.9", color="tab:red"))
-    ax.text(0.05, 4.0, "loop\nR times\n(+step emb\neach pass)", fontsize=7,
-            color="tab:red", ha="center")
+                                "arc3,rad=-1.8", color="tab:red"))
+    ax.text(3.35, 3.5, "loop\nR times\n(+step emb\neach pass)", fontsize=7,
+            color="tab:red", ha="center", va="center")
     ax.text(0.15, 5.6, "input:\n[f(0..63), x, k]", fontsize=7)
     ax.text(2.0, 0.9, "example: 1x4 = 1 shared block, R=4\n"
                       "         4x4 = 4 shared blocks, R=4\n"