pt_loader.js

/* Browser loader for PyTorch .pt checkpoints (torch.save zip format).
 *
 * Parses the zip container (stored entries), unpickles data.pkl with a
 * targeted protocol-2 unpickler (dicts, lists, strings, numbers, tensors
 * via torch._utils._rebuild_tensor_v2 + persistent storage ids), and
 * returns plain JS objects with tensors as {shape, data: Float32Array}.
 *
 * This is not a general unpickler — it covers exactly what torch.save
 * emits for state dicts and simple metadata, and throws loudly otherwise.
 */

function parseZip(buf) {
  const dv = new DataView(buf);
  // find end-of-central-directory
  let eocd = -1;
  for (let i = buf.byteLength - 22; i >= 0; i--) {
    if (dv.getUint32(i, true) === 0x06054b50) { eocd = i; break; }
  }
  if (eocd < 0) throw new Error('not a zip (.pt) file');
  const count = dv.getUint16(eocd + 10, true);
  let off = dv.getUint32(eocd + 16, true);
  const entries = {};
  const td = new TextDecoder();
  for (let e = 0; e < count; e++) {
    if (dv.getUint32(off, true) !== 0x02014b50) throw new Error('bad central dir');
    const method = dv.getUint16(off + 10, true);
    const size = dv.getUint32(off + 24, true);
    const nameLen = dv.getUint16(off + 28, true);
    const extraLen = dv.getUint16(off + 30, true);
    const cmtLen = dv.getUint16(off + 32, true);
    const lho = dv.getUint32(off + 42, true);
    const name = td.decode(new Uint8Array(buf, off + 46, nameLen));
    if (method !== 0) throw new Error('compressed zip entry (unsupported): ' + name);
    // local header: skip its name/extra to find data start
    const lnl = dv.getUint16(lho + 26, true);
    const lel = dv.getUint16(lho + 28, true);
    entries[name] = { offset: lho + 30 + lnl + lel, size };
    off += 46 + nameLen + extraLen + cmtLen;
  }
  return entries;
}

function unpickle(buf, pklOff, pklLen, getStorage) {
  const b = new Uint8Array(buf, pklOff, pklLen);
  const dv = new DataView(buf, pklOff, pklLen);
  const td = new TextDecoder();
  let p = 0;
  const stack = [], marks = [], memo = [];
  const MARKER = Symbol('mark');
  const readLine = () => {
    let s = p; while (b[p] !== 10) p++;
    return td.decode(b.subarray(s, p++));
  };
  const popMark = () => {
    const items = stack.splice(marks.pop());
    return items;
  };
  while (p < b.length) {
    const op = b[p++];
    switch (op) {
      case 0x80: p++; break;                                  // PROTO
      case 0x95: p += 8; break;                               // FRAME (proto 4)
      case 0x28: marks.push(stack.length); break;             // MARK (
      case 0x7d: stack.push({}); break;                       // EMPTY_DICT }
      case 0x5d: stack.push([]); break;                       // EMPTY_LIST ]
      case 0x29: stack.push([]); break;                       // EMPTY_TUPLE )
      case 0x4e: stack.push(null); break;                     // NONE N
      case 0x88: stack.push(true); break;                     // NEWTRUE
      case 0x89: stack.push(false); break;                    // NEWFALSE
      case 0x4b: stack.push(b[p++]); break;                   // BININT1 K
      case 0x4d: stack.push(dv.getUint16(p, true)); p += 2; break;   // BININT2 M
      case 0x4a: stack.push(dv.getInt32(p, true)); p += 4; break;    // BININT J
      case 0x8a: {                                            // LONG1
        const n = b[p++]; let v = 0n;
        for (let i = 0; i < n; i++) v |= BigInt(b[p + i]) << BigInt(8 * i);
        if (n && (b[p + n - 1] & 0x80)) v -= 1n << BigInt(8 * n);
        p += n; stack.push(Number(v)); break;
      }
      case 0x47: {                                            // BINFLOAT G (big-endian)
        stack.push(dv.getFloat64(p, false)); p += 8; break;
      }
      case 0x58: {                                            // BINUNICODE X
        const n = dv.getUint32(p, true); p += 4;
        stack.push(td.decode(b.subarray(p, p + n))); p += n; break;
      }
      case 0x8c: {                                            // SHORT_BINUNICODE
        const n = b[p++];
        stack.push(td.decode(b.subarray(p, p + n))); p += n; break;
      }
      case 0x55: {                                            // SHORT_BINSTRING U
        const n = b[p++];
        stack.push(td.decode(b.subarray(p, p + n))); p += n; break;
      }
      case 0x71: memo[b[p++]] = stack[stack.length - 1]; break;        // BINPUT q
      case 0x72: memo[dv.getUint32(p, true)] = stack[stack.length - 1]; p += 4; break; // LONG_BINPUT
      case 0x94: memo.push(stack[stack.length - 1]); break;            // MEMOIZE
      case 0x68: stack.push(memo[b[p++]]); break;                      // BINGET h
      case 0x6a: stack.push(memo[dv.getUint32(p, true)]); p += 4; break; // LONG_BINGET
      case 0x63: {                                            // GLOBAL c
        const module = readLine(), name = readLine();
        stack.push({ __global__: module + '.' + name });
        break;
      }
      case 0x74: stack.push(popMark()); break;                // TUPLE t
      case 0x85: stack.push([stack.pop()]); break;            // TUPLE1
      case 0x86: { const y = stack.pop(), x = stack.pop(); stack.push([x, y]); break; }
      case 0x87: { const z = stack.pop(), y = stack.pop(), x = stack.pop();
        stack.push([x, y, z]); break; }
      case 0x52: {                                            // REDUCE R
        const args = stack.pop(), fn = stack.pop();
        stack.push(applyReduce(fn, args, getStorage)); break;
      }
      case 0x51: {                                            // BINPERSID Q
        const pid = stack.pop();                              // ('storage', type, key, loc, numel)
        const tname = pid[1] && pid[1].__global__ || '';
        stack.push({ __storage__: pid[2], dtype: tname, numel: pid[4] });
        break;
      }
      case 0x73: {                                            // SETITEM s
        const v = stack.pop(), k = stack.pop();
        stack[stack.length - 1][k] = v; break;
      }
      case 0x75: {                                            // SETITEMS u
        const items = popMark(), d = stack[stack.length - 1];
        for (let i = 0; i < items.length; i += 2) d[items[i]] = items[i + 1];
        break;
      }
      case 0x61: stack[stack.length - 1].push(stack.pop()); break;   // APPEND a
      case 0x65: {                                            // APPENDS e
        const items = popMark();
        stack[stack.length - 1].push(...items); break;
      }
      case 0x81: {                                            // NEWOBJ (cls, args)
        const args = stack.pop(), cls = stack.pop();
        stack.push(applyReduce(cls, args, getStorage)); break;
      }
      case 0x62: {                                            // BUILD b
        const state = stack.pop(), obj = stack[stack.length - 1];
        if (state && typeof state === 'object' && obj && typeof obj === 'object')
          Object.assign(obj, state);
        break;
      }
      case 0x2e: return stack.pop();                          // STOP .
      default:
        throw new Error('unpickle: unsupported opcode 0x' + op.toString(16) + ' at ' + (p - 1));
    }
  }
  throw new Error('unpickle: no STOP');
}

function applyReduce(fn, args, getStorage) {
  const g = fn && fn.__global__;
  if (g === 'collections.OrderedDict') return {};
  if (g === 'torch._utils._rebuild_tensor_v2') {
    const [stor, storageOffset, size, stride] = args;
    if (!/Float/.test(stor.dtype))
      throw new Error('only float32 tensors supported, got ' + stor.dtype);
    const raw = getStorage(stor.__storage__);
    const data = new Float32Array(raw.buffer, raw.byteOffset + storageOffset * 4,
                                  size.reduce((a, x) => a * x, 1));
    return { shape: size, data: Float32Array.from(data) };
  }
  if (g === 'torch.serialization._get_layout') return null;
  throw new Error('unpickle: unsupported constructor ' + g);
}

export async function loadCheckpoint(url) {
  const resp = await fetch(url);
  if (!resp.ok) throw new Error('fetch ' + url + ': ' + resp.status);
  const buf = await resp.arrayBuffer();
  const entries = parseZip(buf);
  const pklName = Object.keys(entries).find(n => n.endsWith('/data.pkl'));
  if (!pklName) throw new Error('no data.pkl in ' + url);
  const prefix = pklName.slice(0, -'data.pkl'.length);
  const getStorage = key => {
    const e = entries[prefix + 'data/' + key];
    if (!e) throw new Error('missing storage ' + key);
    return new Uint8Array(buf, e.offset, e.size);
  };
  const e = entries[pklName];
  return unpickle(buf, e.offset, e.size, getStorage);
}

/* ---------- small MLP helpers (row-major torch Linear weights) ---------- */
export function linear(w, b, x, out) {
  const [rows, cols] = w.shape, wd = w.data, bd = b.data;
  for (let r = 0; r < rows; r++) {
    let s = bd[r];
    const o = r * cols;
    for (let c = 0; c < cols; c++) s += wd[o + c] * x[c];
    out[r] = s;
  }
  return out;
}
export function silu(v) {
  for (let i = 0; i < v.length; i++) v[i] = v[i] / (1 + Math.exp(-v[i]));
  return v;
}
export function relu(v) {
  for (let i = 0; i < v.length; i++) if (v[i] < 0) v[i] = 0;
  return v;
}
export function softplus(v) {
  for (let i = 0; i < v.length; i++) v[i] = Math.log1p(Math.exp(-Math.abs(v[i]))) + Math.max(v[i], 0);
  return v;
}

/* Generic feed-forward runner for a torch nn.Sequential of Linear layers.
 * `layers` lists the state_dict key prefixes in order (e.g. ['net.0',
 * 'net.2','net.4']); `act` is applied after every layer except the last,
 * `outAct` (optional) after the last. Returns forward(inputFloat32)->Float32Array. */
export function buildMLP(sd, layers, act = relu, outAct = null) {
  const W = layers.map(p => sd[p + '.weight']);
  const B = layers.map(p => sd[p + '.bias']);
  const bufs = W.map(w => new Float32Array(w.shape[0]));
  return function forward(x) {
    let h = x;
    for (let i = 0; i < W.length; i++) {
      linear(W[i], B[i], h, bufs[i]);
      if (i < W.length - 1) act(bufs[i]);
      else if (outAct) outAct(bufs[i]);
      h = bufs[i];
    }
    return h;
  };
}
export function argmax(v) {
  let bi = 0, bv = v[0];
  for (let i = 1; i < v.length; i++) if (v[i] > bv) { bv = v[i]; bi = i; }
  return bi;
}

/* Fold the water token (id 5) of the unified projector into a flat weight
 * array for the WGSL MLP, plus the constants the lambda rule needs.
 * Layout (f32): W1(64x6) b1(64) W2(64x64) b2(64) W3(3x64) b3(3)
 *               W4(64x3) b4(64) W5(64x64) b5(64) w6row0(64) b6(1)  */
export function foldWaterToken(ck) {
  const sd = ck.state_dict, TOK = 5;
  const t = sd['tokens.weight'].data.slice(TOK * 4, TOK * 4 + 4);
  const foldTok = (w, b, nKeep) => {
    const [rows, cols] = w.shape, wd = w.data;
    const W = new Float32Array(rows * nKeep), B = Float32Array.from(b.data);
    for (let r = 0; r < rows; r++) {
      for (let c = 0; c < nKeep; c++) W[r * nKeep + c] = wd[r * cols + c];
      for (let c = nKeep; c < cols; c++) B[r] += wd[r * cols + c] * t[c - nKeep];
    }
    return [W, B];
  };
  const [W1, b1] = foldTok(sd['enc.0.weight'], sd['enc.0.bias'], 6);
  const W2 = sd['enc.2.weight'].data, b2 = sd['enc.2.bias'].data;
  const W3 = sd['enc.4.weight'].data, b3 = sd['enc.4.bias'].data;
  const [W4, b4] = foldTok(sd['dec.0.weight'], sd['dec.0.bias'], 3);
  const W5 = sd['dec.2.weight'].data, b5 = sd['dec.2.bias'].data;
  const w6 = sd['dec.4.weight'].data.slice(0, 64);      // row 0 only (lambda slot)
  const b6 = sd['dec.4.bias'].data[0];

  const parts = [W1, b1, W2, b2, W3, b3, W4, b4, W5, b5, w6,
                 Float32Array.of(b6)];
  let total = parts.reduce((a, x) => a + x.length, 0);
  total = Math.ceil(total / 4) * 4;
  const flat = new Float32Array(total);
  let o = 0;
  for (const x of parts) { flat.set(x, o); o += x.length; }

  // r00 = run(0)[0] — the zero-anchor constant, computed on CPU once
  const s = x => x / (1 + Math.exp(-x));
  let h = Array.from(b1, s);
  let h2 = new Array(64);
  for (let r = 0; r < 64; r++) { let a = b2[r];
    for (let c = 0; c < 64; c++) a += W2[r * 64 + c] * h[c]; h2[r] = s(a); }
  let z = new Array(3);
  for (let r = 0; r < 3; r++) { let a = b3[r];
    for (let c = 0; c < 64; c++) a += W3[r * 64 + c] * h2[c]; z[r] = a; }
  let g1 = new Array(64);
  for (let r = 0; r < 64; r++) { let a = b4[r];
    for (let c = 0; c < 3; c++) a += W4[r * 3 + c] * z[c]; g1[r] = s(a); }
  let g2 = new Array(64);
  for (let r = 0; r < 64; r++) { let a = b5[r];
    for (let c = 0; c < 64; c++) a += W5[r * 64 + c] * g1[c]; g2[r] = s(a); }
  let r00 = b6;
  for (let c = 0; c < 64; c++) r00 += w6[c] * g2[c];

  return { weights: flat, r00, scales: ck.fluid_scales, ref: ck.fluid_ref };
}

/* Interior-lattice |grad C|^2 estimate — used to convert g2 between the
 * training discretization and the demo's (lambda ~ 1/g2, so the ratio
 * calibrates units; exact for the rational rule up to the tiny eps). */
export function g2Lattice(h, d, mass, rho0) {
  const spiky = -45 / (Math.PI * Math.pow(h, 6));
  const reach = Math.ceil(h / d);
  let g2 = 0;
  for (let a = -reach; a <= reach; a++)
    for (let b = -reach; b <= reach; b++)
      for (let c = -reach; c <= reach; c++) {
        if (!a && !b && !c) continue;
        const r2 = (a * a + b * b + c * c) * d * d;
        if (r2 >= h * h) continue;
        const r = Math.sqrt(r2);
        const coef = spiky * (h - r) * (h - r) / r * (mass / rho0);
        g2 += coef * coef * r2;
      }
  return g2;
}

/* W6 warm-start net: rotation-equivariant gated combination of 4 history
 * vectors. Returns forward(V, out): V is Float32Array(12), the rows
 * [c1, c2, h*v, h2*a] flattened; writes the predicted vec3 into out. */
export function buildWarmStartNet(sd) {
  const w0 = sd['mlp.0.weight'], b0 = sd['mlp.0.bias'];
  const w2 = sd['mlp.2.weight'], b2 = sd['mlp.2.bias'];
  const w4 = sd['mlp.4.weight'], b4 = sd['mlp.4.bias'];
  const inv = new Float32Array(10), h1 = new Float32Array(32),
        h2 = new Float32Array(32), gains = new Float32Array(4);
  const iu = [[0,0],[0,1],[0,2],[0,3],[1,1],[1,2],[1,3],[2,2],[2,3],[3,3]];
  return function forward(V, out) {   // V: Float32Array(12) rows [c1,c2,hv,h2a]
    let scale2 = 1e-24;
    for (let k = 0; k < 10; k++) {
      const [i, j] = iu[k];
      inv[k] = V[i*3]*V[j*3] + V[i*3+1]*V[j*3+1] + V[i*3+2]*V[j*3+2];
      if (i === j) scale2 += inv[k];
    }
    for (let k = 0; k < 10; k++) inv[k] /= scale2;
    silu(linear(w0, b0, inv, h1));
    silu(linear(w2, b2, h1, h2));
    linear(w4, b4, h2, gains);
    out[0] = out[1] = out[2] = 0;
    for (let i = 0; i < 4; i++) {
      out[0] += gains[i] * V[i*3];
      out[1] += gains[i] * V[i*3+1];
      out[2] += gains[i] * V[i*3+2];
    }
    return out;
  };
}