| """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() |
|
|
| |
| if model is not None: |
| hardest = np.argsort(clues)[:3] |
| 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") |
|
|
| |
| 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") |
|
|
| |
| 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) |
|
|
| |
| 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") |
|
|
