visualize.py

"""
visualize.py -- live visualizer for Wheeler-LM: watch the "wavefunction of the universe"
mixer as the model generates, token by token.

There is no colony and no space here -- the Wheeler-DeWitt block is TIMELESS, so the
right things to watch are the quantities the equation is actually about:

Captures per generated token (via instrument.Recorder, which WheelerDeWittBlock writes to):
  * the TIMELIKE mode Psi_0 of each layer -- the emergent "clock" (volume/scale direction)
  * the Hamiltonian-constraint residual |H| per layer -- how close to the physical H=0 surface
  * the full K minisuperspace modes Psi for one layer -- the state on superspace
  * per-layer block norm + attention + top predictions + token stream

Renders one self-contained HTML dashboard (pure inline SVG, no external scripts / CDN):
  * "Emergent time" -- Psi_0 (the clock) traced across every generated token
  * "Hamiltonian constraint |H|" -- traced across tokens (lower = more physical)
  * "Minisuperspace state" -- the K modes at the current token, mode 0 (timelike) highlighted
  * "|H| per layer" heat, attention, token stream, top predictions, trait activity

Usage:
  python visualize.py --prompt "the universe" --tokens 60
"""
import argparse, json, os, sys, webbrowser
import torch

from model import QuazimotoLM, QuazimotoConfig
import instrument

PKG_DIR = os.path.dirname(os.path.abspath(__file__))


def find_ckpt(path):
    if path and os.path.isfile(path):
        return path
    folder = (os.path.dirname(path) if path else os.path.join(PKG_DIR, "chkpt")) or "."
    if not os.path.isdir(folder):
        return None
    pts = [os.path.join(folder, f) for f in os.listdir(folder) if f.endswith(".pt")]
    return max(pts, key=os.path.getmtime) if pts else None


def load_tokenizer(tok_dir):
    sys.path.insert(0, tok_dir)
    from spike_tokenizer import SpikeTokenizer
    return SpikeTokenizer(vocab_file=os.path.join(tok_dir, "tokenizer.json"))


@torch.no_grad()
def run_capture(model, cfg, tok, ids, n_tokens, temperature, top_k, device):
    pl = cfg.n_layer // 2
    rec = instrument.Recorder(phase_layer=pl, attn_layer=pl)
    instrument.set_rec(rec)
    idx = torch.tensor([ids], device=device)
    try:
        for _ in range(n_tokens):
            cond = idx[:, -cfg.block_size:]
            rec.begin()
            logits, _, _ = model(cond)
            lg = torch.nan_to_num(logits[:, -1, :].float(), nan=0.0,
                                  posinf=1e4, neginf=-1e4) / max(temperature, 1e-6)
            if top_k:
                v = torch.topk(lg, min(top_k, lg.size(-1))).values
                lg = lg.masked_fill(lg < v[:, [-1]], float("-inf"))
            probs = torch.softmax(lg, dim=-1)
            nxt = (torch.argmax(lg, -1, keepdim=True) if temperature <= 1e-3
                   else torch.multinomial(probs, 1))
            top = torch.topk(torch.softmax(logits[:, -1].float(), -1), 5)
            rec.end(token=int(nxt),
                    char=tok.decode([int(nxt)], skip_special_tokens=False),
                    top=[[int(t), round(float(p), 3)] for t, p in zip(top.indices[0], top.values[0])])
            idx = torch.cat([idx, nxt], dim=1)
    finally:
        instrument.set_rec(None)
    return rec.frames, idx[0].tolist()


def main():
    p = argparse.ArgumentParser(description="Wheeler-LM emergent-time / constraint visualizer")
    p.add_argument("--ckpt", default="")
    p.add_argument("--tok_dir", default=PKG_DIR)
    p.add_argument("--prompt", default="the universe")
    p.add_argument("--tokens", type=int, default=60)
    p.add_argument("--temperature", type=float, default=0.0)
    p.add_argument("--top_k", type=int, default=40)
    p.add_argument("--device", default="cpu")
    p.add_argument("--out", default=os.path.join(PKG_DIR, "viz.html"))
    p.add_argument("--no_open", action="store_true")
    args = p.parse_args()

    path = find_ckpt(args.ckpt)
    if path is None:
        print("No checkpoint found; pass --ckpt or train one."); return
    ckpt = torch.load(path, map_location=args.device, weights_only=False)
    cfg = QuazimotoConfig(**ckpt["family_config"])
    model = QuazimotoLM(cfg); model.load_state_dict(ckpt["model"], strict=False)
    model.to(args.device).eval()
    tok = load_tokenizer(args.tok_dir)
    print(f"loaded {path} (step {ckpt.get('step')}) | capturing {args.tokens} tokens ...")

    ids = tok.encode(args.prompt, add_special_tokens=False)
    frames, _ = run_capture(model, cfg, tok, ids, args.tokens,
                            args.temperature, args.top_k or None, args.device)

    data = {
        "meta": {"ckpt": os.path.basename(path), "step": ckpt.get("step"),
                 "n_layer": cfg.n_layer, "wdw_modes": cfg.wdw_modes,
                 "wdw_steps": cfg.wdw_steps, "phase_layer": cfg.n_layer // 2,
                 "prompt": args.prompt},
        "prompt_chars": [tok.decode([t], skip_special_tokens=False) for t in ids],
        "frames": frames,
    }
    html = HTML_TEMPLATE.replace("/*__DATA__*/", json.dumps(data))
    with open(args.out, "w", encoding="utf-8") as f:
        f.write(html)
    print(f"wrote {args.out} ({os.path.getsize(args.out)//1024} KB)")
    if not args.no_open:
        webbrowser.open("file://" + os.path.abspath(args.out))


HTML_TEMPLATE = r"""<!doctype html>
<html><head><meta charset="utf-8"><title>Wheeler-LM live</title>
<style>
  :root{--bg:#0d1117;--panel:#161b22;--ink:#e6edf3;--mut:#8b949e;--ac:#58a6ff;--hot:#f0883e;--spore:#3fb950;--red:#f85149}
  *{box-sizing:border-box} body{margin:0;background:var(--bg);color:var(--ink);font:13px/1.4 ui-monospace,Menlo,Consolas,monospace}
  header{padding:10px 16px;background:var(--panel);border-bottom:1px solid #30363d;display:flex;gap:16px;align-items:center;flex-wrap:wrap}
  h1{font-size:15px;margin:0;color:var(--ac)} .mut{color:var(--mut)}
  #wrap{display:grid;grid-template-columns:320px 1fr 300px;gap:12px;padding:12px}
  .panel{background:var(--panel);border:1px solid #30363d;border-radius:8px;padding:10px}
  .panel h2{font-size:12px;margin:0 0 8px;color:var(--mut);text-transform:uppercase;letter-spacing:.5px}
  #stream{line-height:1.9;max-height:120px;overflow:auto}
  .tk{padding:1px 2px;border-radius:3px;cursor:pointer;white-space:pre}
  .tk.prompt{color:var(--mut)} .tk.cur{background:var(--ac);color:#000} .tk:hover{background:#30363d}
  .ctrl{display:flex;gap:10px;align-items:center} input[type=range]{width:260px}
  button{background:#21262d;color:var(--ink);border:1px solid #30363d;border-radius:6px;padding:4px 10px;cursor:pointer}
  button:hover{border-color:var(--ac)} svg{display:block;width:100%;height:auto}
  .bar{height:10px;background:#21262d;border-radius:5px;overflow:hidden;margin:2px 0}.bar>div{height:100%}
  table{border-collapse:collapse;width:100%}td{padding:1px 4px}.lbl{color:var(--mut);width:52px}
  .grid{display:grid;gap:3px}.cell{height:15px;border-radius:2px}.small{font-size:11px;color:var(--mut)}
</style></head><body>
<header><h1>Wheeler-LM &middot; wavefunction of the universe</h1><span class="mut" id="meta"></span>
  <span style="flex:1"></span>
  <div class="ctrl"><button id="play">▶ play</button><input type="range" id="seek" min="0" value="0"><span id="pos" class="mut"></span></div>
</header>
<div id="wrap">
  <div style="display:flex;flex-direction:column;gap:12px">
    <div class="panel"><h2>Token stream</h2><div id="stream"></div></div>
    <div class="panel"><h2>Trait activity (this token)</h2><div id="traits"></div></div>
    <div class="panel"><h2>Top predictions</h2><div id="top"></div></div>
  </div>
  <div style="display:flex;flex-direction:column;gap:12px">
    <div class="panel"><h2>Emergent time &mdash; timelike mode &Psi;&#8320; (layer <span id="pl"></span>) across tokens</h2><div id="clock"></div></div>
    <div class="panel"><h2>Hamiltonian constraint |H| across tokens (lower = more physical, H&Psi;=0)</h2><div id="hcon"></div></div>
    <div class="panel"><h2>Minisuperspace state &mdash; K modes &Psi; at this token (orange = timelike)</h2><div id="modes"></div></div>
  </div>
  <div style="display:flex;flex-direction:column;gap:12px">
    <div class="panel"><h2>|H| per layer (this token)</h2><div id="heat"></div></div>
    <div class="panel"><h2>Attention &mdash; layer <span id="al"></span> (last token &rarr; context)</h2><div id="attn"></div></div>
    <div class="panel"><h2>Notes</h2><div class="small" id="notes"></div></div>
  </div>
</div>
<script>
const D=/*__DATA__*/; const M=D.meta,F=D.frames,NL=M.n_layer,K=M.wdw_modes,PL=M.phase_layer;
let i=0,playing=false,timer=null; const $=id=>document.getElementById(id);
$("meta").textContent=`${M.ckpt} · step ${M.step} · ${K} modes · ${M.wdw_steps} wave steps · ${NL} layers · "${M.prompt}"`;
$("pl").textContent=PL; $("al").textContent=PL;
const seek=$("seek"); seek.max=F.length-1;
const hue2=v=>`hsl(${(1-Math.max(0,Math.min(1,v)))*20+0},80%,${30+35*Math.max(0,Math.min(1,v))}%)`; // red hot heat

// per-token series from the phase layer: Psi_0 (clock) and |H| (constraint residual)
const clock=F.map(f=>f.rings&&f.rings[PL]?f.rings[PL].psi[0]:0);
const hcon =F.map(f=>f.rings&&f.rings[PL]?Math.abs(f.rings[PL].R[0]):0);

// generic SVG line chart of a series, with a cursor at index i and a zero baseline
function lineChart(series,color,W,H,cur,zero){
  const n=series.length; if(!n) return "";
  let mn=Math.min(...series),mx=Math.max(...series); if(zero){mn=Math.min(mn,0);mx=Math.max(mx,0);}
  if(mx-mn<1e-9)mx=mn+1; const pad=6;
  const X=k=>pad+(W-2*pad)*(n<2?0.5:k/(n-1));
  const Y=v=>pad+(H-2*pad)*(1-(v-mn)/(mx-mn));
  let pts=series.map((v,k)=>`${X(k).toFixed(1)},${Y(v).toFixed(1)}`).join(" ");
  let z = zero?`<line x1="${pad}" y1="${Y(0).toFixed(1)}" x2="${W-pad}" y2="${Y(0).toFixed(1)}" stroke="#30363d" stroke-dasharray="3 3"/>`:"";
  let cx=X(cur).toFixed(1);
  let cy=Y(series[cur]).toFixed(1);
  return `<svg viewBox="0 0 ${W} ${H}">${z}
    <polyline points="${pts}" fill="none" stroke="${color}" stroke-width="1.6"/>
    <line x1="${cx}" y1="${pad}" x2="${cx}" y2="${H-pad}" stroke="${color}" stroke-opacity="0.4"/>
    <circle cx="${cx}" cy="${cy}" r="3.2" fill="${color}"/>
    <text x="${pad}" y="12" fill="#8b949e" font-size="10">${mx.toFixed(3)}</text>
    <text x="${pad}" y="${H-3}" fill="#8b949e" font-size="10">${mn.toFixed(3)}</text></svg>`;
}

// K minisuperspace modes as up/down bars around a zero baseline; mode 0 (timelike) hot
function modeBars(theta){
  const n=theta.length,W=560,H=140,pad=8; if(!n)return "";
  const amx=Math.max(1e-6,...theta.map(Math.abs));
  const bw=(W-2*pad)/n, mid=H/2;
  let bars=theta.map((v,k)=>{const hgt=(H/2-pad)*(Math.abs(v)/amx);
    const y=v>=0?mid-hgt:mid; const col=(k===0)?"var(--hot)":"var(--ac)";
    return `<rect x="${(pad+k*bw).toFixed(1)}" y="${y.toFixed(1)}" width="${Math.max(1,bw-1).toFixed(1)}" height="${hgt.toFixed(1)}" fill="${col}"/>`;}).join("");
  return `<svg viewBox="0 0 ${W} ${H}"><line x1="${pad}" y1="${mid}" x2="${W-pad}" y2="${mid}" stroke="#30363d"/>${bars}
    <text x="${pad}" y="12" fill="#8b949e" font-size="10">timelike &Psi;&#8320; = ${theta[0].toFixed(3)}</text></svg>`;
}

const stream=$("stream");
D.prompt_chars.forEach(c=>{const s=document.createElement("span");s.className="tk prompt";s.textContent=esc(c);stream.appendChild(s);});
F.forEach((f,k)=>{const s=document.createElement("span");s.className="tk gen";s.textContent=esc(f.char);s.onclick=()=>{i=k;render()};s.dataset.k=k;stream.appendChild(s);});
function esc(c){return c.replace(/\n/g,"⏎").replace(/ /g,"·");}

function render(){
  const f=F[i]; $("pos").textContent=`${i+1}/${F.length}`; seek.value=i;
  document.querySelectorAll(".tk.gen").forEach(s=>s.classList.toggle("cur",+s.dataset.k===i));
  $("clock").innerHTML=lineChart(clock,"#f0883e",560,150,i,true);
  $("hcon").innerHTML =lineChart(hcon, "#f85149",560,110,i,false);
  $("modes").innerHTML=(f.phases&&f.phases.theta)?modeBars(f.phases.theta):"<span class='small'>n/a</span>";
  // |H| per layer heat row
  const hL=(f.rings||[]).map(r=>Math.abs(r.R[0])); const hmx=Math.max(...hL,1e-6);
  let hh=`<div class="grid" style="grid-template-columns:repeat(${NL},1fr)">`;
  for(let l=0;l<NL;l++){const v=hL[l]||0;hh+=`<div class="cell" title="L${l} |H|=${v.toFixed(4)}" style="background:${hue2(v/hmx)}"></div>`;}
  $("heat").innerHTML=hh+"</div><div class='small'>bright red = far from H=0</div>";
  // trait activity (block norms + HRM/MoE)
  const qn=f.quaz_norm||[],mx=Math.max(...qn,1e-6);let L=qn.map((_,l)=>"wave L"+l),V=qn.map(v=>v/mx);
  if(f.traits&&f.traits.hrm!=null){L.push("HRM");V.push(Math.min(1,f.traits.hrm/mx));}
  if(f.traits&&f.traits.moe!=null){L.push("MoE");V.push(Math.min(1,f.traits.moe/mx));}
  let th="<table>";V.forEach((v,k)=>{th+=`<tr><td class="lbl">${L[k]}</td><td style="width:100%"><div class="bar"><div style="width:${(v*100).toFixed(0)}%;background:${hue2(v)}"></div></div></td></tr>`;});
  $("traits").innerHTML=th+"</table>";
  let tp="<table>";f.top.forEach(([t,p])=>{tp+=`<tr><td class="lbl">${p.toFixed(2)}</td><td><div class="bar"><div style="width:${(p*100).toFixed(0)}%;background:var(--ac)"></div></div></td><td class="small">id ${t}</td></tr>`;});
  $("top").innerHTML=tp+"</table>";
  if(f.attn){const w=f.attn.w,m=Math.max(...w,1e-6);let h="<div style='display:flex;flex-wrap:wrap;gap:1px'>";
    w.forEach((a,p)=>{h+=`<div title="pos ${p}: ${a.toFixed(3)}" style="width:8px;height:14px;background:${hue2(a/m)}"></div>`;});
    $("attn").innerHTML=h+`</div><div class='small'>${w.length} context positions</div>`;}
  else $("attn").innerHTML="<span class='small'>n/a</span>";
  $("notes").innerHTML=`Modes: ${K} · wave steps: ${M.wdw_steps}<br>&Psi;&#8320; is the volume/scale (timelike) mode &mdash; the emergent clock; there is no external time.<br>|H| is the Hamiltonian-constraint residual; training drives it toward 0 (H&Psi;=0, a physical state).<br>Orange bar = timelike mode; blue = spacelike. Click any token or use the slider.`;
}
seek.oninput=()=>{i=+seek.value;render()};
$("play").onclick=()=>{playing=!playing;$("play").textContent=playing?"⏸ pause":"▶ play";if(playing)timer=setInterval(()=>{i=(i+1)%F.length;render();},400);else clearInterval(timer);};
document.onkeydown=e=>{if(e.key==="ArrowRight"){i=Math.min(F.length-1,i+1);render();}if(e.key==="ArrowLeft"){i=Math.max(0,i-1);render();}};
render();
</script></body></html>"""


if __name__ == "__main__":
    main()