Spaces:
Running
Running
| // In-browser port of brain.py + the step() of server.py: leaky integrate-and-fire network | |
| // over the FlyWire v783 connectome (Shiu et al. 2024 constants), 1 ms steps. | |
| const V0 = -52, VTH = -45, T_MBR = 20, TAU = 5, DT = 1, REFRAC = 2; | |
| const DECAY = Math.exp(-DT / TAU), K_V = DT / T_MBR; | |
| let meta, N, indptr, idx, w, c; // connectome (shared, read-only) | |
| let v, g, rfc, ring, tick = 0; // state of this simulation | |
| let groupBits, groupNames, isDN, dnPos, recent, drive = [], spkFlag, forcedBuf; | |
| // mushroom-body learning (real KC -> MBON synapses, depressed when a dopamine neuron is co-active) | |
| let mb = null, kcOrd, mbonOrd, danOrd, plOff, plK, plPre, plPost, mult, kcTrace, danTrace, mbCnt; | |
| const ETA = 0.04; | |
| let RECOVER = 0.00025, clock = 1; // memory fades on the fly-day clock (set by the page) | |
| async function gunzip(buf) { | |
| if (buf[0] === 0x1f && buf[1] === 0x8b) { | |
| const s = new Blob([buf]).stream().pipeThrough(new DecompressionStream('gzip')); | |
| return await new Response(s).arrayBuffer(); | |
| } | |
| return buf.buffer; | |
| } | |
| async function load(url, rawBytes) { | |
| const res = await fetch(url); | |
| const len = +res.headers.get('Content-Length') || 0; | |
| const reader = res.body.getReader(); | |
| const parts = []; let got = 0; | |
| for (;;) { | |
| const { done, value } = await reader.read(); | |
| if (done) break; | |
| parts.push(value); got += value.length; | |
| if (len) postMessage({ type: 'progress', file: url, frac: got / len }); | |
| } | |
| const buf = new Uint8Array(got); let o = 0; | |
| for (const p of parts) { buf.set(p, o); o += p.length; } | |
| if (rawBytes) return buf.buffer; // a part of a larger gzip stream | |
| return await gunzip(buf); | |
| } | |
| async function init(ver) { | |
| meta = await (await fetch('meta.json?v=' + ver)).json(); | |
| N = meta.n; c = meta.wsyn * meta.gain; | |
| postMessage({ type: 'status', text: 'downloading connectome (~37 MB)…' }); | |
| indptr = new Uint32Array(await load('indptr.dat')); | |
| const parts = []; // delta.dat is split in 4 parts of one gzip stream | |
| for (let i = 0; i < 4; i++) parts.push(new Uint8Array(await load(`delta${i}.dat`, true))); | |
| let tot = 0; for (const p of parts) tot += p.length; | |
| const joined = new Uint8Array(tot); { let o = 0; for (const p of parts) { joined.set(p, o); o += p.length; } } | |
| idx = new Uint32Array(await gunzip(joined)); | |
| w = new Int16Array(await load('weight.dat')); | |
| postMessage({ type: 'status', text: 'decoding…' }); | |
| for (let r = 0; r < N; r++) { // undo the per-row delta coding | |
| let acc = 0; | |
| for (let k = indptr[r], e = indptr[r + 1]; k < e; k++) { acc = k === indptr[r] ? idx[k] : acc + idx[k]; idx[k] = acc; } | |
| } | |
| groupNames = Object.keys(meta.groups); | |
| groupBits = new Uint8Array(N); | |
| groupNames.forEach((gn, b) => { for (const i of meta.groups[gn]) groupBits[i] |= 1 << b; }); | |
| isDN = new Uint8Array(N); dnPos = new Int32Array(N).fill(-1); | |
| meta.dn.forEach((i, p) => { isDN[i] = 1; dnPos[i] = p; }); | |
| recent = new Float32Array(meta.dn.length); | |
| spkFlag = new Uint8Array(N); | |
| if (meta.mb) setupLearning(); | |
| reset(); | |
| postMessage({ type: 'ready', info: { neurons: N, synapses: meta.synapses, connections: meta.connections, | |
| pos: meta.pos, dn: meta.dn } }); | |
| } | |
| function setupLearning() { | |
| mb = meta.mb; | |
| kcOrd = new Int32Array(N).fill(-1); mbonOrd = new Int32Array(N).fill(-1); danOrd = new Int32Array(N).fill(-1); | |
| mb.kc.forEach((i, o) => { kcOrd[i] = o; }); | |
| mb.mbon.forEach((i, o) => { mbonOrd[i] = o; }); | |
| mb.pam.forEach((i, o) => { danOrd[i] = o; }); | |
| mb.ppl.forEach((i, o) => { danOrd[i] = mb.pam.length + o; }); | |
| const K = [], pre = [], post = []; plOff = new Int32Array(mb.kc.length + 1); | |
| mb.kc.forEach((i, o) => { | |
| plOff[o] = K.length; | |
| for (let k = indptr[i], e = indptr[i + 1]; k < e; k++) { | |
| const m = mbonOrd[idx[k]]; | |
| if (m >= 0) { K.push(k); pre.push(i); post.push(m); } | |
| } | |
| }); | |
| plOff[mb.kc.length] = K.length; | |
| plK = Int32Array.from(K); plPre = Int32Array.from(pre); plPost = Int32Array.from(post); | |
| mult = new Float32Array(K.length).fill(1); | |
| kcTrace = new Float32Array(mb.kc.length); | |
| danTrace = new Float32Array(mb.pam.length + mb.ppl.length); | |
| mbCnt = new Float32Array(mb.mbon.length); | |
| } | |
| function reset() { | |
| v = new Float32Array(N).fill(V0); g = new Float32Array(N); rfc = new Uint8Array(N); | |
| ring = [new Float32Array(N), new Float32Array(N)]; tick = 0; | |
| if (recent) recent.fill(0); | |
| if (mb) { mult.fill(1); kcTrace.fill(0); danTrace.fill(0); mbCnt.fill(0); } | |
| } | |
| // one 1 ms step; returns the list of neurons that spiked | |
| function step() { | |
| const cur = ring[tick++ & 1]; | |
| // forced (sensory) spikes, decided on the refractory state before this step | |
| let nf = 0; | |
| for (let j = 0; j < drive.length; j++) { | |
| const d = drive[j], ids = d.ids, p = d.p; | |
| for (let k = 0; k < ids.length; k++) { | |
| const i = ids[k]; | |
| if (rfc[i] <= 0 && Math.random() < p && !spkFlag[i]) { spkFlag[i] = 1; forcedBuf[nf++] = i; } | |
| } | |
| } | |
| const spikes = []; | |
| for (let i = 0; i < N; i++) { | |
| const gi = g[i]; | |
| const vi = v[i] + (V0 - v[i] + gi) * K_V; | |
| g[i] = gi * DECAY + cur[i]; | |
| cur[i] = 0; | |
| v[i] = vi; | |
| const r = rfc[i]; | |
| if (spkFlag[i] || (vi > VTH && r <= 0)) spikes.push(i); | |
| if (r > 0) rfc[i] = r - 1; | |
| } | |
| for (let k = 0; k < nf; k++) spkFlag[forcedBuf[k]] = 0; | |
| for (let s = 0; s < spikes.length; s++) { | |
| const i = spikes[s]; | |
| v[i] = V0; g[i] = 0; rfc[i] = REFRAC; | |
| for (let k = indptr[i], e = indptr[i + 1]; k < e; k++) cur[idx[k]] += w[k] * c; | |
| if (mb) { | |
| const o = kcOrd[i]; | |
| if (o >= 0) { | |
| kcTrace[o] += 1; | |
| for (let p = plOff[o], pe = plOff[o + 1]; p < pe; p++) cur[idx[plK[p]]] += w[plK[p]] * c * (mult[p] - 1); | |
| } else { | |
| const d = danOrd[i]; if (d >= 0) danTrace[d] += 1; | |
| const m = mbonOrd[i]; if (m >= 0) mbCnt[m] += 1; | |
| } | |
| } | |
| } | |
| return spikes; | |
| } | |
| function runStep(stim) { | |
| const t0 = performance.now(); | |
| drive = []; let nd = 0; | |
| for (const [name, list] of Object.entries(meta.channels)) { | |
| const s = Math.min(1, Math.max(0, +stim[name] || 0)); | |
| if (s > 0.001 && list.length) { drive.push({ ids: list, p: s * meta.max_rate }); nd += list.length; } | |
| } | |
| if (!forcedBuf || forcedBuf.length < nd) forcedBuf = new Int32Array(Math.max(nd, 1024)); | |
| const G = groupNames.length, cnt = new Array(G).fill(0); | |
| const steps = meta.steps_per_call; | |
| let total = 0; const shown = []; | |
| for (let s = 0; s < steps; s++) { | |
| const sp = step(); | |
| total += sp.length; | |
| for (let k = 0; k < sp.length; k++) { | |
| const i = sp[k], b = groupBits[i]; | |
| if (b) for (let q = 0; q < G; q++) if (b & (1 << q)) cnt[q]++; | |
| if (isDN[i]) recent[dnPos[i]]++; | |
| } | |
| if (sp.length <= 120) for (let k = 0; k < sp.length; k++) shown.push(sp[k]); | |
| else for (let k = 0; k < 120; k++) shown.push(sp[(Math.random() * sp.length) | 0]); | |
| } | |
| for (let p = 0; p < recent.length; p++) recent[p] *= 0.9; | |
| let mbA = 0, mbV = 0, depressed = 0; | |
| if (mb) { | |
| const nP = mb.pam.length, nM = mb.mbon.length; | |
| // dopamine drive onto each MBON = share of its own dopamine input that is active | |
| const mod = new Float32Array(nM); | |
| for (let m = 0; m < nM; m++) { | |
| let num = 0, tot = 0; | |
| for (let d = 0; d < nP; d++) { const wgt = mb.Mp[d][m]; tot += wgt; num += Math.min(1, danTrace[d] / 3) * wgt; } | |
| for (let d = 0; d < mb.ppl.length; d++) { const wgt = mb.Ml[d][m]; tot += wgt; num += Math.min(1, danTrace[nP + d] / 3) * wgt; } | |
| mod[m] = Math.min(1, Math.max(0, num / (tot + 1e-9) - 0.6) / 0.4); | |
| } | |
| for (let p = 0; p < mult.length; p++) { | |
| const kt = Math.min(1, kcTrace[kcOrd[plPre[p]]] / 2), dm = mod[plPost[p]]; | |
| let v = mult[p]; | |
| if (kt > 0 && dm > 0) v *= 1 - ETA * kt * dm; | |
| v += (1 - v) * RECOVER * clock; | |
| mult[p] = v < 0.1 ? 0.1 : v; | |
| if (v < 0.9) depressed++; | |
| } | |
| for (let m = 0; m < nM; m++) { if (mb.avoid[m]) mbV += mbCnt[m]; else mbA += mbCnt[m]; } | |
| kcTrace.forEach((_, i) => { kcTrace[i] *= 0.5; }); | |
| for (let d = 0; d < danTrace.length; d++) danTrace[d] *= 0.6; | |
| mbCnt.fill(0); | |
| } | |
| while (shown.length > 1500) shown.splice((Math.random() * shown.length) | 0, 1); | |
| const out = {}; | |
| groupNames.forEach((gn, q) => { out[gn] = Math.min(1, cnt[q] / steps / meta.ref[gn]); }); | |
| const order = Array.from(recent.keys()).sort((a, b) => recent[b] - recent[a]).slice(0, 5); | |
| out.spikes_per_ms = total / steps; | |
| out.spiking = shown; | |
| out.top_dn = order.filter(p => recent[p] > 0.5).map(p => [meta.dn_names[p], Math.round(recent[p] * 10) / 10]); | |
| out.mb_approach = mbA; out.mb_avoid = mbV; | |
| out.mb_base_approach = mb ? mb.base_approach * steps / 500 : 0; | |
| out.mb_base_avoid = mb ? mb.base_avoid * steps / 500 : 0; | |
| out.depressed = mb ? depressed / Math.max(1, mult.length) : 0; | |
| out.compute_ms = Math.round(performance.now() - t0); | |
| return out; | |
| } | |
| onmessage = (e) => { | |
| const m = e.data; | |
| if (m.type === 'init') init(m.v).catch(err => postMessage({ type: 'error', text: String(err) })); | |
| else if (m.type === 'clock') clock = Math.max(0.1, +m.speed || 1); | |
| else if (m.type === 'reset') reset(); | |
| else if (m.type === 'getMemory') { | |
| const q = new Uint8Array(mult ? mult.length : 0); | |
| for (let p = 0; p < q.length; p++) q[p] = Math.round(mult[p] * 255); | |
| postMessage({ type: 'memory', data: q }, [q.buffer]); | |
| } else if (m.type === 'setMemory' && mult && m.data && m.data.length === mult.length) { | |
| for (let p = 0; p < mult.length; p++) mult[p] = Math.max(0.1, m.data[p] / 255); | |
| } else if (m.type === 'forget' && mult) mult.fill(1); | |
| else if (m.type === 'step') postMessage({ type: 'out', id: m.id, out: runStep(m.stim || {}) }); | |
| }; | |