FIG-001 Figure

"""FIG-001: DiT architecture, left half (paper Fig. 2 left).

Self-contained generator for knowledge card FIG-001 (SYSTEM.md D19):
solve -> render -> overwrite FIG-001.svg next to this script.
Run: kblayout/.venv/bin/python figures/FIG-001.py
"""
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "kblayout"))

from kblayout.core.figure import Figure
from kblayout.ext import layout
from kblayout.ext.line import edge
from kblayout.ext.rect import rect
from kblayout.ext.text import text_block, measure as text_block_measure

FIGURE = Figure(id="dit-latent")


def node(nid, label, cx, cy, *,
         font_size: float = 14.0, gap: float = 10.0):
    """Rect wraps the text by one ring of `gap` (composed, not
    hardcoded: the text metric drives the rect size - measure() is
    the single source of the font geometry, D8-mu)."""
    tw, th = text_block_measure(label, font_size, pad=0.0)
    w, h = tw + 2 * gap, th + 2 * gap
    r = rect(FIGURE, nid, w, h, cx=cx, cy=cy)
    text_block(FIGURE, r, label, pad=0.0)
    return r


# ---- main chain (upward); Patchify column is offset LEFT so that
# Patchify and Embed sit side-by-side under the (widened) block and
# BOTH go straight up from their North anchors into its bottom
# side - as in the paper figure ----
n_nz = node("noised_latent", "Noised\nLatent", -45, 0)
n_pch = node("patchify", "Patchify", -45, -110)
n_b1 = node("block_1", "DiT Block\n× N", 0, -230, gap=29.0)
n_norm = node("norm", "Layer Norm", 0, -330)
n_lr = node("linear_reshape", "Linear and\nReshape", 0, -430)

# ---- outputs: two parallel prediction heads ----
n_noise = node("noise", "Noise", -90, -530)
n_sigma = node("sigma", "Σ", 90, -530)

# ---- conditioning side column (t/y below Embed, arrows up, as in
# the paper: both feed the block from its bottom side) ----
n_t = node("timestep", "Timestep t", 60, -20)
n_y = node("label", "Label y", 170, -20)
n_em = node("embed", "Embed", 45, -110)

# ---- edges: main chain ----
# pch->block and em->block use to_body: the contact point on the
# block bottom is NOT declared (it lands wherever the normal entry
# meets the energy); declaring b1.bottom (bottom-mid anchor x=0)
# forced a diagonal seed which axis attraction folded into a
# stair-step zigzag (上-横-上). Free normal contact is the paper
# figure's arrow semantics.
main_chain = [n_nz, n_pch, n_b1, n_norm, n_lr]
# explicit chain edges (no auto-for: pch->norm would cross the block)
edge(FIGURE, n_nz.top, n_pch.bottom, bends=0, U0=(0, -1), U1=(0, -1))
edge(FIGURE, n_pch.top, to_body=n_b1, bends=0, sign=-1.0)
edge(FIGURE, n_b1.top, n_norm.bottom, bends=0, U0=(0, -1), U1=(0, -1))
edge(FIGURE, n_norm.top, n_lr.bottom, bends=0, U0=(0, -1), U1=(0, -1))
# outputs split from Linear and Reshape top
edge(FIGURE, n_lr.top, n_noise.bottom, bends=0, U0=(0, -1), U1=(0, -1))
edge(FIGURE, n_lr.top, n_sigma.bottom, bends=0, U0=(0, -1), U1=(0, -1))

# ---- edges: conditioning column ----
edge(FIGURE, n_t.top, n_em.bottom, bends=0, U0=(0, -1), U1=(0, -1))
edge(FIGURE, n_y.top, n_em.bottom, bends=0, U0=(0, -1), U1=(0, -1))


# ---- explicit layout relations ----
for src, dst in zip(main_chain, main_chain[1:]):
    FIGURE.constraints.append(layout.flow_progress(src=src, dst=dst, gap=64.0))
edge(FIGURE, n_em.top, to_body=n_b1, bends=0, sign=-1.0)
FIGURE.constraints.append(
    layout.align_x(refs=[n_nz.center, n_pch.center], name="layout.pch_col"))
FIGURE.constraints.append(
    layout.align_x(refs=[n_b1.center, n_norm.center, n_lr.center],
                   name="layout.main_col"))
FIGURE.constraints.append(
    layout.align_x(refs=[n_t.center, n_em.center], name="layout.em_col"))
FIGURE.constraints.append(
    layout.align_y(refs=[n_t.center, n_y.center], name="layout.io_row"))
FIGURE.constraints.append(
    layout.align_y(refs=[n_noise.bottom, n_sigma.bottom],
                   name="layout.output_row"))
for a, b in [(n_nz, n_y), (n_t, n_y), (n_b1, n_y), (n_lr, n_y),
             (n_pch, n_em)]:
    FIGURE.constraints.append(layout.pair_no_overlap(a=a, b=b, gap=16.0))
# experiment 2026-09-04: single absolute pin (chain entry) - the
# other boxes are free to drift; only relative terms couple them
FIGURE.constraints.append(
    layout.origin_anchor(refs=[n_nz.center]))


if __name__ == "__main__":
    from kblayout.kgrad.solver import solve_kgrad as solve
    from kblayout.render.svg import render
    from kblayout.report import backfill_card, print_report

    # latest engine: whole-solve loop in one Mojo call (A5: requires
    # KGRAD_MOJO_LIB pointing at kblayout/kgrad-mojo; no silent fallback)
    res = solve(FIGURE, steps=3000, lr=0.1, momentum=0.9, backend="mojo")
    svg = render(res, FIGURE)
    out = Path(__file__).with_suffix(".svg")
    out.write_text(svg, encoding="utf-8")
    print(f"{out.name}: loss {float(res.loss_final):.2f}")
    print_report(res)
    print(backfill_card(Path(__file__).with_suffix(".md"), res, svg))