File size: 8,624 Bytes
8c85090 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 | """Sudoku figures that show the model actually solving, plus difficulty curves.
Everything is drawn from real runs and real test puzzles - the grids below are
genuine model output on held-out instances, not illustrations.
"""
import glob, json, os, sys
from collections import defaultdict
import numpy as np
import torch
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from model import RecurrentDenoiser, Transformer
import diffusion as dfn
from sudoku import CELLS, GRID, SudokuTokenizer
RUNS = os.environ.get("RUNS", "runs")
OUT = os.environ.get("OUT", "figures")
HARD = os.environ.get("HARD", "data/sudoku_hard")
os.makedirs(OUT, exist_ok=True)
dev = "cuda" if torch.cuda.is_available() else "cpu"
tok = SudokuTokenizer()
P = np.load(os.path.join(HARD, "test_puz.npy"))
S = np.load(os.path.join(HARD, "test_sol.npy"))
clues = (P > 0).sum(1)
def draw_grid(ax, grid, given=None, wrong=None, title="", fs=9):
"""Render one 9x9 board. Given clues bold, wrong cells shaded red."""
g = np.asarray(grid).reshape(GRID, GRID)
gv = np.asarray(given).reshape(GRID, GRID) if given is not None else None
wr = np.asarray(wrong).reshape(GRID, GRID) if wrong is not None else None
ax.set_xlim(0, 9); ax.set_ylim(0, 9); ax.invert_yaxis()
ax.set_xticks([]); ax.set_yticks([])
for i in range(GRID):
for j in range(GRID):
if wr is not None and wr[i, j]:
ax.add_patch(Rectangle((j, i), 1, 1, color="#f5b7b1", zorder=0))
elif gv is not None and gv[i, j]:
ax.add_patch(Rectangle((j, i), 1, 1, color="#eaeded", zorder=0))
v = int(g[i, j])
if v > 0:
bold = gv is not None and gv[i, j]
ax.text(j + .5, i + .55, str(v), ha="center", va="center",
fontsize=fs, fontweight="bold" if bold else "normal",
color="#000" if bold else ("#922b21" if (wr is not None and wr[i, j])
else "#1a5276"), zorder=2)
for k in range(10):
lw = 1.9 if k % 3 == 0 else 0.5
ax.plot([k, k], [0, 9], color="k", lw=lw)
ax.plot([0, 9], [k, k], color="k", lw=lw)
ax.set_title(title, fontsize=8.5)
def load_best():
"""The strongest trained model available."""
best, bacc = None, -1
for d in sorted(glob.glob(f"{RUNS}/rec1-*")) + sorted(glob.glob(f"{RUNS}/inj-rec*")):
rf, mf = os.path.join(d, "result.json"), os.path.join(d, "model.pt")
if not (os.path.exists(rf) and os.path.exists(mf)):
continue
r = json.load(open(rf))
acc = max(v["board"] for v in r["decode"].values())
if acc > bacc:
bacc, best = acc, (d, r["args"])
if best is None:
return None, None
d, cfg = best
m = RecurrentDenoiser(len(tok), cfg["d"], cfg["layers"], cfg["heads"], cfg["pe"],
max_len=CELLS + 8,
recurrences=cfg.get("R_infer") or cfg["R"]).to(dev)
m.load_state_dict(torch.load(os.path.join(d, "model.pt"), map_location=dev))
m.eval()
return m, cfg
model, cfg = load_best()
# ---- Figure 5: the model solving, pass by pass ------------------------------
if model is not None:
hardest = np.argsort(clues)[:3] # fewest clues = hardest
budgets = [1, 4, 16, 61]
fig, axes = plt.subplots(len(hardest), len(budgets) + 2,
figsize=(2.0 * (len(budgets) + 2), 2.05 * len(hardest)))
for r, idx in enumerate(hardest):
puz, sol = P[idx], S[idx]
given = puz > 0
draw_grid(axes[r][0], puz, given,
title=f"puzzle ({given.sum()} clues)" if r == 0 else f"({given.sum()} clues)")
for c, K in enumerate(budgets):
with torch.no_grad():
gold = torch.from_numpy(sol - 1)[None].to(dev)
bl = torch.from_numpy(~given)[None].to(dev)
x = torch.where(bl, torch.full_like(gold, tok.mask), gold)
y, _ = dfn.fixed_k_decode(model, x, tok, dev, K, fillable=bl)
pred = y[0].cpu().numpy() + 1
wrong = (pred != sol) & (~given)
ok = not wrong.any()
draw_grid(axes[r][c + 1], pred, given, wrong,
title=(f"{K} pass{'es' if K > 1 else ''}" if r == 0 else "")
+ (" ✓" if ok else f" {wrong.sum()} wrong"))
draw_grid(axes[r][-1], sol, given, title="solution" if r == 0 else "")
fig.suptitle("Recurrent masked diffusion solving held-out hard Sudoku "
"(grey = given clue, red = wrong cell)", fontsize=10)
fig.tight_layout(rect=[0, 0, 1, 0.96])
fig.savefig(f"{OUT}/sud_fig5_grids.png", dpi=165)
print("wrote sud_fig5_grids.png")
# ---- Figure 6: how the board fills across recurrences -----------------------
if model is not None:
idx = int(np.argsort(clues)[0])
puz, sol = P[idx], S[idx]
given = puz > 0
Rs = [1, 2, 4, 8, 16, 32, 64]
fig, axes = plt.subplots(1, len(Rs), figsize=(2.0 * len(Rs), 2.35))
for c, R in enumerate(Rs):
model.recurrences = R
with torch.no_grad():
gold = torch.from_numpy(sol - 1)[None].to(dev)
bl = torch.from_numpy(~given)[None].to(dev)
x = torch.where(bl, torch.full_like(gold, tok.mask), gold)
y, _ = dfn.fixed_k_decode(model, x, tok, dev, 16, fillable=bl)
pred = y[0].cpu().numpy() + 1
wrong = (pred != sol) & (~given)
draw_grid(axes[c], pred, given, wrong,
title=f"R={R} {(~wrong[~given]).mean()*100:.0f}% cells", fs=7.5)
model.recurrences = cfg.get("R_infer") or cfg["R"]
fig.suptitle("The same block applied more times: iterative refinement, "
"trained at R=32", fontsize=10)
fig.tight_layout(rect=[0, 0, 1, 0.93])
fig.savefig(f"{OUT}/sud_fig6_recurrence.png", dpi=165)
print("wrote sud_fig6_recurrence.png")
# ---- Figure 7: accuracy against puzzle difficulty ---------------------------
if model is not None:
model.recurrences = cfg.get("R_infer") or cfg["R"]
preds = []
with torch.no_grad():
for i in range(0, len(P), 200):
gold = torch.from_numpy(S[i:i + 200] - 1).to(dev)
bl = torch.from_numpy(P[i:i + 200] == 0).to(dev)
x = torch.where(bl, torch.full_like(gold, tok.mask), gold)
y, _ = dfn.fixed_k_decode(model, x, tok, dev, 61, fillable=bl)
preds.append(y.cpu().numpy() + 1)
pred = np.concatenate(preds)
solved = (pred == S).all(axis=1)
ks = sorted({int(c) for c in clues if (clues == c).sum() >= 15})
acc = [solved[clues == k].mean() * 100 for k in ks]
nn = [(clues == k).sum() for k in ks]
fig, ax = plt.subplots(figsize=(6.8, 4.2))
ax.plot(ks, acc, "o-", color="#1a5276", lw=2, ms=7, label="recurrent diffusion (ours)")
ax.axhline(3.0, ls=":", color="#c0392b", lw=1.4)
ax.text(ks[0], 5.5, "constraint propagation, whole set: 3.0%", fontsize=7.5,
color="#c0392b")
for k, a, n in zip(ks, acc, nn):
ax.annotate(f"n={n}", (k, a), fontsize=6.5, xytext=(0, -13),
textcoords="offset points", ha="center", color="#555")
ax.set_xlabel("clues given (fewer = harder)")
ax.set_ylabel("board accuracy (%)")
ax.set_title("Accuracy against puzzle difficulty, held-out test set")
ax.grid(alpha=.3); ax.legend(fontsize=8); ax.set_ylim(-3, 105)
ax.invert_xaxis()
fig.tight_layout(); fig.savefig(f"{OUT}/sud_fig7_difficulty.png", dpi=170)
print("wrote sud_fig7_difficulty.png")
json.dump({"by_clue": {int(k): float(a) for k, a in zip(ks, acc)},
"n_by_clue": {int(k): int(n) for k, n in zip(ks, nn)}},
open(os.path.join(OUT, "sud_by_clue.json"), "w"), indent=2)
# ---- Figure 8: benchmark difficulty, ours vs the saturated standard split ---
fig, ax = plt.subplots(figsize=(6.6, 4.0))
ax.bar(["SATNet split\n(~36 clues)", "ours\n(20–30 clues)"], [100.0, 3.0],
color=["#aab7b8", "#c0392b"], width=.55)
ax.set_ylabel("solved by constraint propagation alone (%)")
ax.set_title("Why we generated a new benchmark:\nthe standard split is saturated")
for i, v in enumerate([100.0, 3.0]):
ax.text(i, v + 2.5, f"{v:.1f}%", ha="center", fontweight="bold")
ax.set_ylim(0, 112); ax.grid(alpha=.3, axis="y")
fig.tight_layout(); fig.savefig(f"{OUT}/sud_fig8_saturation.png", dpi=170)
print("wrote sud_fig8_saturation.png")
|