visualize.py

"""
visualize.py -- live visualizer for Mycel-LM: watch the fungal colony grow in 3D as
the model generates, token by token.

Captures per generated token (via instrument.Recorder, which MycelBlock writes to):
  * tip POSITIONS (all pos_dim axes) for one layer's colony -- the 3D scatter
  * per-tip growth TRAIL over the growth steps -- the hyphae extending this token
  * per-layer colony density + trait-station sites + attention + token stream

Renders one self-contained HTML dashboard (data embedded -> opens via file://). The
colony is a Three.js network you can drag-rotate / scroll-zoom: hyphal tips linked to
their nearest neighbours as filaments (so it reads as a mycelial web, not loose dots),
tips coloured by local density, with trait STATIONS as orange wire-spheres and faint
grey growth trails. Three.js loads from a CDN (needs network the first time).

Usage:
  python visualize.py --prompt "the mycelium spreads" --tokens 50
"""
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):
    rec = instrument.Recorder(phase_layer=cfg.n_layer // 2, attn_layer=cfg.n_layer // 2)
    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="Mycel-LM 3D colony visualizer")
    p.add_argument("--ckpt", default="")
    p.add_argument("--tok_dir", default=PKG_DIR)
    p.add_argument("--prompt", default="the mycelium spreads")
    p.add_argument("--tokens", type=int, default=50)
    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)

    pl = cfg.n_layer // 2
    blk = model.layers[pl].quaz
    stations = blk.stations.anchors.detach().cpu().tolist() if getattr(blk, "use_stations", False) else []

    data = {
        "meta": {"ckpt": os.path.basename(path), "step": ckpt.get("step"),
                 "n_layer": cfg.n_layer, "n_tips": cfg.n_tips, "pos_dim": cfg.mycel_pos_dim,
                 "bound": cfg.osc_bound, "phase_layer": pl, "prompt": args.prompt,
                 "stations": stations,
                 "tip_rings": bool(getattr(cfg, "use_tip_rings", False)),
                 "tip_ring_size": getattr(cfg, "tip_ring_size", 0)},
        "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>Mycel-LM live</title>
<script src="https://cdnjs.cloudflare.com/ajax/libs/three.js/r128/three.min.js"></script>
<script src="https://cdn.jsdelivr.net/npm/three@0.128.0/examples/js/controls/OrbitControls.js"></script>
<style>
  :root{--bg:#0d1117;--panel:#161b22;--ink:#e6edf3;--mut:#8b949e;--ac:#58a6ff;--hot:#f0883e;--spore:#3fb950}
  *{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(--spore)} .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(--spore);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(--spore)} svg{display:block;width:100%;height:auto}
  #colony3d{width:100%;height:440px;border-radius:6px;overflow:hidden}
  .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>Mycel-LM &middot; colony 3D</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>Colony &mdash; layer <span id="pl"></span> (drag rotate, scroll zoom; wire-spheres = trait stations)</h2>
      <div id="colony3d"></div></div>
    <div class="panel"><h2>Colony density per layer (bright = dense/clustered)</h2><div id="heat"></div></div>
  </div>
  <div style="display:flex;flex-direction:column;gap:12px">
    <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,NT=M.n_tips,PD=M.pos_dim,BND=M.bound;
let i=0,playing=false,timer=null; const $=id=>document.getElementById(id);
$("meta").textContent=`${M.ckpt} · step ${M.step} · ${NT} tips · ${PD}D colony · ${NL} layers · "${M.prompt}"`;
$("pl").textContent=M.phase_layer; $("al").textContent=M.phase_layer;
const seek=$("seek"); seek.max=F.length-1;
const hue2=v=>`hsl(${(1-Math.max(0,Math.min(1,v)))*140},70%,${30+45*Math.max(0,Math.min(1,v))}%)`;

// hsl(green->orange) -> [r,g,b] in 0..1 for Three.js vertex colours (density mode)
function colRGB(v){v=Math.max(0,Math.min(1,v));const h=(1-v)*140/360,s=0.7,l=0.32+0.42*v;
  const a=s*Math.min(l,1-l),f=n=>{const k=(n+h*12)%12;return l-a*Math.max(-1,Math.min(k-3,9-k,1));};return[f(0),f(8),f(4)];}
// ring PHASE -> [r,g,b]: hue = phase around the wheel, saturation/brightness = coherence r.
// Synchronized tips share a hue (colony turns one colour); desynced tips are a rainbow.
function phaseRGB(psi,r){const h=(psi/(2*Math.PI))+0.5,s=0.25+0.65*Math.max(0,Math.min(1,r)),l=0.5;
  const a=s*Math.min(l,1-l),f=n=>{const k=(n+h*12)%12;return l-a*Math.max(-1,Math.min(k-3,9-k,1));};return[f(0),f(8),f(4)];}

// ---- Three.js colony ----
const KNN=3;                       // links per tip -> the mycelial web
const MAXE=NT*KNN;                 // max filament segments
let scene,camera,renderer,controls,geom,points,webGeom,web,trailGeom,trail;
function init3d(){
  const el=$("colony3d"),W=el.clientWidth||600,H=440;
  scene=new THREE.Scene(); scene.background=new THREE.Color(0x0d1117);
  camera=new THREE.PerspectiveCamera(55,W/H,0.1,1000); camera.position.set(BND*2.4,BND*1.7,BND*2.4);
  renderer=new THREE.WebGLRenderer({antialias:true}); renderer.setSize(W,H); el.appendChild(renderer.domElement);
  controls=new THREE.OrbitControls(camera,renderer.domElement); controls.enableDamping=true; controls.target.set(0,0,0);
  scene.add(new THREE.LineSegments(new THREE.EdgesGeometry(new THREE.BoxGeometry(2*BND,2*BND,2*BND)),
    new THREE.LineBasicMaterial({color:0x30363d})));
  // filament web: each tip linked to its nearest neighbours -> looks like hyphae, not dots
  webGeom=new THREE.BufferGeometry();
  webGeom.setAttribute('position',new THREE.BufferAttribute(new Float32Array(MAXE*2*3),3));
  webGeom.setAttribute('color',new THREE.BufferAttribute(new Float32Array(MAXE*2*3),3));
  web=new THREE.LineSegments(webGeom,new THREE.LineBasicMaterial({vertexColors:true,transparent:true,opacity:0.55}));
  scene.add(web);
  // growth trails: each tip's path over the growth steps (the hypha extending this token)
  const TS=(F[0]&&F[0].phases&&F[0].phases.traj)?F[0].phases.traj.length:0;
  if(TS>1){trailGeom=new THREE.BufferGeometry();
    trailGeom.setAttribute('position',new THREE.BufferAttribute(new Float32Array(NT*(TS-1)*2*3),3));
    trail=new THREE.LineSegments(trailGeom,new THREE.LineBasicMaterial({color:0x8b949e,transparent:true,opacity:0.35}));
    scene.add(trail);}
  geom=new THREE.BufferGeometry();
  geom.setAttribute('position',new THREE.BufferAttribute(new Float32Array(NT*3),3));
  geom.setAttribute('color',new THREE.BufferAttribute(new Float32Array(NT*3),3));
  points=new THREE.Points(geom,new THREE.PointsMaterial({size:BND*0.07,vertexColors:true}));
  scene.add(points);
  (M.stations||[]).forEach(s=>{const mesh=new THREE.Mesh(new THREE.SphereGeometry(BND*0.06,10,10),
    new THREE.MeshBasicMaterial({color:0xf0883e,wireframe:true}));
    mesh.position.set(s[0]||0,s[1]||0,s[2]||0); scene.add(mesh);});
  window.addEventListener('resize',()=>{const w=el.clientWidth||600;camera.aspect=w/H;camera.updateProjectionMatrix();renderer.setSize(w,H);});
  (function loop(){requestAnimationFrame(loop);controls.update();renderer.render(scene,camera);})();
}
function updateColony(f){
  if(!geom||!f.phases)return;
  const th=f.phases.theta,tp=f.phases.tip_psi,tr=f.phases.tip_r;
  const pos=geom.attributes.position.array,col=geom.attributes.color.array;
  const P=[]; for(let t=0;t<NT;t++)P.push([th[t*PD],th[t*PD+1],PD>2?th[t*PD+2]:0]);
  let colr;
  if(tp){                                  // colour by ring phase (synchronization view)
    colr=P.map((_,t)=>phaseRGB(tp[t],tr?tr[t]:0.6));
  }else{                                   // colour by local density (no rings)
    const r2=(BND*0.25)**2, dens=P.map(a=>{let d=0;P.forEach(b=>{const dx=a[0]-b[0],dy=a[1]-b[1],dz=a[2]-b[2];if(dx*dx+dy*dy+dz*dz<r2)d++;});return d;});
    const mx=Math.max(...dens,1); colr=dens.map(d=>colRGB(d/mx));
  }
  for(let t=0;t<NT;t++){pos[t*3]=P[t][0];pos[t*3+1]=P[t][1];pos[t*3+2]=P[t][2];
    const c=colr[t];col[t*3]=c[0];col[t*3+1]=c[1];col[t*3+2]=c[2];}
  geom.attributes.position.needsUpdate=true; geom.attributes.color.needsUpdate=true;
  // rebuild the filament web: connect each tip to its KNN nearest neighbours
  if(webGeom){const wp=webGeom.attributes.position.array,wc=webGeom.attributes.color.array;let e=0;
    for(let a=0;a<NT&&e<MAXE;a++){
      const nb=[]; for(let b=0;b<NT;b++){if(b===a)continue;
        const dx=P[a][0]-P[b][0],dy=P[a][1]-P[b][1],dz=P[a][2]-P[b][2];nb.push([dx*dx+dy*dy+dz*dz,b]);}
      nb.sort((x,y)=>x[0]-y[0]);
      for(let n=0;n<Math.min(KNN,nb.length)&&e<MAXE;n++){const b=nb[n][1],ca=colr[a],cb=colr[b];
        wp[e*6]=P[a][0];wp[e*6+1]=P[a][1];wp[e*6+2]=P[a][2];
        wp[e*6+3]=P[b][0];wp[e*6+4]=P[b][1];wp[e*6+5]=P[b][2];
        wc[e*6]=ca[0];wc[e*6+1]=ca[1];wc[e*6+2]=ca[2];wc[e*6+3]=cb[0];wc[e*6+4]=cb[1];wc[e*6+5]=cb[2];e++;}}
    webGeom.setDrawRange(0,e*2);webGeom.attributes.position.needsUpdate=true;webGeom.attributes.color.needsUpdate=true;}
  // growth trails: draw each tip's path across the captured growth steps
  if(trail&&f.phases.traj){const tj=f.phases.traj,TS=tj.length,tp2=trailGeom.attributes.position.array;let g=0;
    for(let t=0;t<NT;t++)for(let s=0;s<TS-1;s++){
      const A=tj[s],Bp=tj[s+1];
      tp2[g++]=A[t*PD];tp2[g++]=A[t*PD+1];tp2[g++]=PD>2?A[t*PD+2]:0;
      tp2[g++]=Bp[t*PD];tp2[g++]=Bp[t*PD+1];tp2[g++]=PD>2?Bp[t*PD+2]:0;}
    trailGeom.setDrawRange(0,NT*(TS-1)*2);trailGeom.attributes.position.needsUpdate=true;}
}

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 heat(f){
  let h=`<div class="grid" style="grid-template-columns:auto repeat(${NL},1fr)"><div class="small"></div>`;
  for(let l=0;l<NL;l++)h+=`<div class="small" style="text-align:center">L${l}</div>`;
  h+=`<div class="small">density</div>`; const mx=Math.max(...f.rings.map(r=>r.R[0]),1e-6);
  for(let l=0;l<NL;l++){const s=f.rings[l]?f.rings[l].R[0]:0;h+=`<div class="cell" title="L${l} ${s.toFixed(3)}" style="background:${hue2(s/mx)}"></div>`;}
  $("heat").innerHTML=h+"</div>";
}
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));
  updateColony(f); heat(f);
  const qn=f.quaz_norm,mx=Math.max(...qn,1e-6);let L=qn.map((_,l)=>"grow L"+l),V=qn.map(v=>v/mx);
  if(f.traits.hrm!=null){L.push("HRM");V.push(Math.min(1,f.traits.hrm/mx));}
  if(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=`Tips: ${NT} · stations: ${(M.stations||[]).length}<br>Green = sparse growing edge, orange = dense core.<br>Filaments link each tip to its nearest neighbours (the mycelial web); faint grey trails are each tip's growth path this token.<br>Orange wire-spheres = trait stations. Drag to orbit, scroll to zoom.`;
}
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();}};
if(window.THREE){init3d();} else {$("colony3d").innerHTML="<div class='small' style='padding:20px'>Three.js failed to load (needs network for the CDN).</div>";}
render();
</script></body></html>"""


if __name__ == "__main__":
    main()