318 lines
17 KiB
TypeScript
318 lines
17 KiB
TypeScript
import assert from 'node:assert/strict';
|
|
import { existsSync, readFileSync } from 'node:fs';
|
|
import { join } from 'node:path';
|
|
import test from 'node:test';
|
|
import { fileURLToPath } from 'node:url';
|
|
import { gunzipSync } from 'node:zlib';
|
|
import type { BrainDataset, BrainMetadata } from '../src/dataset/format';
|
|
import { DEFAULT_LIF_CONFIG, FlyBrain, LifNetwork, MAX_RATE_ROLES, NEURAL_KERNEL_VERSION, ROLE_MASK_WORDS, kernelVersion } from '../src/model/lif';
|
|
import { DEFAULT_RETINA_CONFIG, projectFrame } from '../src/model/retina';
|
|
import { Xorshift32 } from '../src/model/rng';
|
|
import { toyDataset, xorshift } from './fixtures/toy-dataset';
|
|
import { FlyBrain as LegacyFlyBrain } from './legacy/lif';
|
|
import type { Plasticity as LegacyPlasticity } from './legacy/plasticity';
|
|
import type { RewardModulatedStdp } from '../src/model/plasticity';
|
|
|
|
/** Toy connectome with a single retina column on neuron 0, as the original warm-up test used. */
|
|
function visualToyDataset(): BrainDataset {
|
|
const data = toyDataset();
|
|
data.visualIndices = Uint32Array.of(0);
|
|
data.visualHemisphere = Uint8Array.of(0);
|
|
data.visualXY = Float32Array.of(0, 0);
|
|
return data;
|
|
}
|
|
|
|
/** Deterministic RGBA frame. */
|
|
function randomFrame(width: number, height: number, seed: number): Uint8Array {
|
|
const random = xorshift(seed);
|
|
const rgba = new Uint8Array(width * height * 4);
|
|
for (let i = 0; i < rgba.length; i++) rgba[i] = Math.floor(random() * 256);
|
|
return rgba;
|
|
}
|
|
|
|
/** Gray frame of 4x4 blocks, so nearest-neighbour rounding differences stay inside a block. */
|
|
function blockFrame(width: number, height: number): Uint8Array {
|
|
const rgba = new Uint8Array(width * height * 4);
|
|
for (let y = 0; y < height; y++) {
|
|
for (let x = 0; x < width; x++) {
|
|
const value = (Math.floor(x * 4 / width) * 4 + Math.floor(y * 4 / height)) * 15 + 3;
|
|
const offset = (y * width + x) * 4;
|
|
rgba[offset] = value; rgba[offset + 1] = value; rgba[offset + 2] = value; rgba[offset + 3] = 255;
|
|
}
|
|
}
|
|
return rgba;
|
|
}
|
|
|
|
/** One causal pre(0)->post(1) pairing 10 ms apart, as in the original unit tests. */
|
|
function causal(p: RewardModulatedStdp | LegacyPlasticity, ms: number) {
|
|
p.observe(Uint32Array.of(1), 1, Float64Array.of(ms - 10, -1e6, -1e6, -1e6), ms);
|
|
}
|
|
|
|
test('Xorshift32 reproduces the original kernel stream bit-exactly', () => {
|
|
let legacy = 22_222;
|
|
const legacyUint = () => { let value = legacy; value ^= value << 13; value ^= value >>> 17; value ^= value << 5; legacy = value; return value >>> 0; };
|
|
const rng = new Xorshift32();
|
|
assert.equal(rng.state, 22_222);
|
|
for (let i = 0; i < 10_000; i++) {
|
|
assert.equal(rng.nextUint(), legacyUint());
|
|
assert.equal(rng.state, legacy);
|
|
}
|
|
const checkpoint = rng.state;
|
|
const draws = [rng.next(), rng.next(), rng.next()];
|
|
rng.state = checkpoint;
|
|
assert.deepEqual([rng.next(), rng.next(), rng.next()], draws);
|
|
assert.ok(draws.every(value => value >= 0 && value < 1));
|
|
assert.notEqual(new Xorshift32(7).nextUint(), new Xorshift32().nextUint());
|
|
});
|
|
|
|
test('projectFrame matches the original 160x144 retina projection', () => {
|
|
const xy = Float32Array.of(120, -40, 300, 900, -55, 12, 4, 4, 1000, 1000, 0, 0, 640, 480, 17, 900);
|
|
const hemisphere = Uint8Array.of(0, 1, 0, 1, 1, 0, 1, 0);
|
|
const data = toyDataset();
|
|
data.visualIndices = Uint32Array.of(0, 1, 2, 3, 0, 1, 2, 3);
|
|
data.visualHemisphere = hemisphere;
|
|
data.visualXY = xy;
|
|
const frame = randomFrame(DEFAULT_RETINA_CONFIG.width, DEFAULT_RETINA_CONFIG.height, 3);
|
|
|
|
const legacy = new LegacyFlyBrain(data);
|
|
legacy.setVisualFrame(frame);
|
|
const expected = legacy.exportState().visualDrive;
|
|
|
|
const out = new Float32Array(8);
|
|
projectFrame(frame, DEFAULT_RETINA_CONFIG.width, DEFAULT_RETINA_CONFIG.height,
|
|
{ xy, hemisphere, count: 8 }, DEFAULT_RETINA_CONFIG.gain, out);
|
|
assert.deepEqual(out, expected);
|
|
assert.ok(out.some(value => value !== out[0]), 'frame must actually vary across columns');
|
|
|
|
const brain = new LifNetwork(data);
|
|
brain.setVisualFrame(frame);
|
|
assert.deepEqual(brain.exportState().visualDrive, expected);
|
|
});
|
|
|
|
test('retina projection is resolution independent: 320x288 maps to the same relative pixels', () => {
|
|
const xy = Float32Array.of(0, 0, 0.1, 0.35, 0.6, 0.85, 1, 1, 0.35, 0.6, 0.85, 0.1);
|
|
const hemisphere = Uint8Array.of(1, 1, 0, 1, 0, 1);
|
|
const columns = { xy, hemisphere, count: 6 };
|
|
const small = new Float32Array(6);
|
|
const large = new Float32Array(6);
|
|
projectFrame(blockFrame(160, 144), 160, 144, columns, DEFAULT_RETINA_CONFIG.gain, small);
|
|
projectFrame(blockFrame(320, 288), 320, 288, columns, DEFAULT_RETINA_CONFIG.gain, large);
|
|
assert.deepEqual(large, small);
|
|
assert.equal(new Set(small).size, 6, 'columns must land in six distinct blocks');
|
|
|
|
// Gain scales the drive linearly and the frame size is a per-call override on the network.
|
|
const doubled = new Float32Array(6);
|
|
projectFrame(blockFrame(320, 288), 320, 288, columns, DEFAULT_RETINA_CONFIG.gain * 2, doubled);
|
|
for (let i = 0; i < 6; i++) assert.equal(doubled[i], Math.fround(small[i] * 2));
|
|
|
|
const data = toyDataset();
|
|
data.visualIndices = Uint32Array.of(0, 1, 2, 3, 0, 1);
|
|
data.visualHemisphere = hemisphere;
|
|
data.visualXY = xy;
|
|
const brain = new LifNetwork(data);
|
|
brain.setVisualFrame(blockFrame(320, 288), 320, 288);
|
|
assert.deepEqual(brain.exportState().visualDrive, small);
|
|
});
|
|
|
|
test('LifNetwork is bit-exact with the legacy FlyBrain over 3000 ms on the toy dataset', () => {
|
|
const legacy = new LegacyFlyBrain(visualToyDataset());
|
|
const next = new LifNetwork(visualToyDataset());
|
|
const compare = (label: string) => {
|
|
assert.deepEqual(next.exportState(), legacy.exportState(), label);
|
|
assert.deepEqual(next.rates, legacy.rates, label);
|
|
assert.equal(next.populationRate, legacy.populationRate, label);
|
|
assert.equal(next.ms, legacy.ms, label);
|
|
assert.deepEqual(next.membrane, legacy.membrane, label);
|
|
assert.deepEqual(next.baseline, legacy.baseline, label);
|
|
assert.deepEqual(next.refractory, legacy.refractory, label);
|
|
assert.deepEqual(next.roleNames, legacy.roleNames, label);
|
|
assert.deepEqual(next.plasticity.statistics(), legacy.plasticity.statistics(), label);
|
|
};
|
|
compare('construction');
|
|
assert.equal(legacy.step(100), next.step(100));
|
|
compare('after 100 ms');
|
|
|
|
for (let round = 0; round < 6; round++) {
|
|
const frame = randomFrame(160, 144, round + 1);
|
|
legacy.setVisualFrame(frame); next.setVisualFrame(frame);
|
|
assert.equal(legacy.step(300), next.step(300), `round ${round} visual drive`);
|
|
if (round % 2 === 0) { legacy.reward(120); next.stimulate(120); }
|
|
assert.equal(legacy.step(100), next.step(100), `round ${round} stimulation`);
|
|
causal(legacy.plasticity, legacy.ms); causal(next.plasticity, next.ms);
|
|
const reward = round % 2 === 0 ? -0.25 : 0.5;
|
|
legacy.plasticity.reinforce(reward, legacy.ms); next.plasticity.reinforce(reward, next.ms);
|
|
compare(`round ${round}`);
|
|
}
|
|
|
|
assert.equal(legacy.step(500), next.step(500));
|
|
compare('after 3000 ms');
|
|
assert.equal(next.ms, 3000);
|
|
assert.ok(next.populationRate > 0, 'the network must actually spike');
|
|
assert.ok(next.plasticity.statistics().changed > 0, 'reinforcement must actually move gains');
|
|
assert.ok(next.rates.command_0 >= 0);
|
|
|
|
// A fresh network resumes the legacy checkpoint exactly.
|
|
const resumed = new LifNetwork(visualToyDataset());
|
|
resumed.importState(legacy.exportState());
|
|
assert.deepEqual(resumed.exportState(), legacy.exportState());
|
|
resumed.step(250); legacy.step(250);
|
|
assert.deepEqual(resumed.exportState(), legacy.exportState());
|
|
});
|
|
|
|
test('reward is an alias of stimulate and drives the configured role', () => {
|
|
assert.equal(FlyBrain, LifNetwork);
|
|
const stimulated = new LifNetwork(toyDataset());
|
|
const rewarded = new LifNetwork(toyDataset());
|
|
stimulated.stimulate(50); rewarded.reward(50);
|
|
stimulated.step(10); rewarded.step(10);
|
|
assert.deepEqual(rewarded.exportState(), stimulated.exportState());
|
|
assert.equal(stimulated.exportState().rewardRemaining, 40);
|
|
// Overlapping pulses take the maximum remainder.
|
|
stimulated.stimulate(10); assert.equal(stimulated.exportState().rewardRemaining, 40);
|
|
stimulated.stimulate(200); assert.equal(stimulated.exportState().rewardRemaining, 200);
|
|
|
|
// The default role is absent from the toy dataset, so stimulation is a no-op there; a
|
|
// configured role that exists moves the membrane.
|
|
const driven = new LifNetwork(toyDataset(), { stimulation: { role: 'motor', drive: 0.5 }, noiseKicks: 0, baselineMax: 0 });
|
|
driven.stimulate(1); driven.step(1);
|
|
assert.ok(driven.membrane[3] > 0.4 && driven.membrane[0] === 0);
|
|
assert.equal(driven.exportState().rewardRemaining, 0);
|
|
});
|
|
|
|
test('default configuration keeps the historical kernel version; numeric changes derive a new one', () => {
|
|
assert.equal(kernelVersion(), NEURAL_KERNEL_VERSION);
|
|
assert.equal(kernelVersion({ ...DEFAULT_LIF_CONFIG }), NEURAL_KERNEL_VERSION);
|
|
assert.equal(new LifNetwork(toyDataset()).version, NEURAL_KERNEL_VERSION);
|
|
const versions = new Set<string>();
|
|
const patches: Partial<typeof DEFAULT_LIF_CONFIG>[] = [
|
|
{ decayMs: 25 }, { threshold: 0.9 }, { refractoryMs: 3 }, { synapseScale: 0.006 }, { baselineMax: 0.05 },
|
|
{ noiseKicks: 200 }, { noiseAmount: 0.5 }, { rateAlpha: 1 / 50 }, { membraneFloor: -3 }, { seed: 1 },
|
|
{ stimulation: { role: 'reward_pam', drive: 0.3 } },
|
|
{ retina: { gain: 0.3, width: 160, height: 144 } },
|
|
{ retina: { gain: 0.20, width: 320, height: 144 } },
|
|
{ retina: { gain: 0.20, width: 160, height: 288 } },
|
|
];
|
|
for (const patch of patches) {
|
|
const version = kernelVersion(patch);
|
|
assert.notEqual(version, NEURAL_KERNEL_VERSION);
|
|
assert.match(version, /^lif-1ms-f64-v2:[0-9a-f]{8}$/);
|
|
versions.add(version);
|
|
}
|
|
assert.equal(versions.size, patches.length, 'each numeric parameter must reach the version hash');
|
|
// Role names are dataset labels, not kernel constants: they never move the version.
|
|
assert.equal(kernelVersion({ stimulation: { role: 'other', drive: 0.20 }, rateRoles: ['command_0'] }), NEURAL_KERNEL_VERSION);
|
|
assert.equal(new LifNetwork(toyDataset(), { decayMs: 25 }).version, kernelVersion({ decayMs: 25 }));
|
|
});
|
|
|
|
test('non-default constants change the dynamics, not only the version string', () => {
|
|
// Weak noise keeps the four-neuron toy network near threshold; saturated drive would hide the
|
|
// threshold and baseline parameters, which only matter while spiking is not already maximal.
|
|
const base = { noiseKicks: 2, noiseAmount: 0.2 };
|
|
const reference = new LifNetwork(visualToyDataset(), base);
|
|
reference.step(600);
|
|
for (const patch of [{ decayMs: 25 }, { threshold: 0.5 }, { refractoryMs: 5 }, { synapseScale: 0.05 },
|
|
{ baselineMax: 0.2 }, { noiseKicks: 6 }, { noiseAmount: 0.4 }, { rateAlpha: 1 / 5 }, { seed: 99 }]) {
|
|
const variant = new LifNetwork(visualToyDataset(), { ...base, ...patch });
|
|
variant.step(600);
|
|
assert.notDeepEqual(variant.exportState(), reference.exportState(), JSON.stringify(patch));
|
|
}
|
|
// membraneFloor only bites on inhibition, so compare the clamped neuron directly.
|
|
const shallow = new LifNetwork(visualToyDataset(), { membraneFloor: -0.01, noiseKicks: 0, baselineMax: 0 });
|
|
const deep = new LifNetwork(visualToyDataset(), { membraneFloor: -50, noiseKicks: 0, baselineMax: 0 });
|
|
for (const brain of [shallow, deep]) { brain.membrane[0] = 10; brain.step(2); }
|
|
assert.ok(shallow.membrane[2] > deep.membrane[2]);
|
|
});
|
|
|
|
test('rate roles default to the dataset order and are capped at the bitmask width', () => {
|
|
const data = toyDataset();
|
|
data.meta.roles = { mbon: [1, 2], command_3: [3], steer_left: [0], kenyon: [0], command_0: [3], proboscis: [1] };
|
|
assert.deepEqual(new LifNetwork(data).roleNames, ['command_3', 'steer_left', 'command_0', 'proboscis']);
|
|
// A `macro_<type>` population is tracked by default like a `command_*` button
|
|
// (`docs/design/macros.md` section 12: a macro is a button), in the dataset's own order.
|
|
const withMacros = toyDataset();
|
|
withMacros.meta.roles = { macro_talk: [1], command_0: [3], kenyon: [0], macro_go_out: [2] };
|
|
assert.deepEqual(new LifNetwork(withMacros).roleNames, ['macro_talk', 'command_0', 'macro_go_out']);
|
|
// An explicit list keeps its own order and drops names the dataset does not declare.
|
|
assert.deepEqual(new LifNetwork(data, { rateRoles: ['mbon', 'absent', 'command_0'] }).roleNames, ['mbon', 'command_0']);
|
|
assert.deepEqual(new LifNetwork(data, { rateRoles: [] }).rates, {});
|
|
|
|
const wide = toyDataset();
|
|
wide.meta.roles = {};
|
|
for (let i = 0; i < MAX_RATE_ROLES; i++) wide.meta.roles[`command_${i}`] = [i % 4];
|
|
const packed = new LifNetwork(wide);
|
|
assert.equal(packed.roleNames.length, MAX_RATE_ROLES);
|
|
assert.ok(packed.roleMasks instanceof Uint32Array);
|
|
assert.equal(packed.roleMasks.length, wide.meta.neurons * ROLE_MASK_WORDS);
|
|
// The top role lives in the last word's top bit, which is what the second word bought
|
|
// (`docs/design/macros.md` section 11 needs thirty-six roles, not fourteen).
|
|
assert.notEqual(packed.roleMasks[3 * ROLE_MASK_WORDS + ROLE_MASK_WORDS - 1] & (1 << 31), 0, 'the last role must occupy the top bit');
|
|
packed.step(200);
|
|
assert.ok(Object.values(packed.rates).every(Number.isFinite));
|
|
assert.ok(packed.rates[`command_${MAX_RATE_ROLES - 1}`] > 0);
|
|
|
|
wide.meta.roles[`command_${MAX_RATE_ROLES}`] = [0];
|
|
assert.throws(() => new LifNetwork(wide), new RegExp(String(MAX_RATE_ROLES)));
|
|
});
|
|
|
|
// --- Oracle on the real FAFB v783 artifacts, when the dataset is present in the worktree.
|
|
|
|
function loadRealDataset(base: string): BrainDataset {
|
|
const meta = JSON.parse(readFileSync(join(base, 'meta.json'), 'utf8')) as BrainMetadata;
|
|
const circuits = JSON.parse(readFileSync(join(base, 'circuit-roles.json'), 'utf8')) as { neurons: number; roles: Record<string, number[]> };
|
|
if (circuits.neurons !== meta.neurons) throw new Error('Circuit roles do not match connectome');
|
|
meta.roles = { ...meta.roles, ...circuits.roles };
|
|
const bytes = (name: string) => new Uint8Array(gunzipSync(readFileSync(join(base, name))));
|
|
const data: BrainDataset = {
|
|
meta,
|
|
indptr: new Uint32Array(bytes('indptr.binz').buffer),
|
|
targets: new Uint32Array(bytes('targets.binz').buffer),
|
|
weights: new Int16Array(bytes('weights.binz').buffer),
|
|
visualIndices: new Uint32Array(bytes('visual-indices.binz').buffer),
|
|
visualHemisphere: bytes('visual-hemisphere.binz'),
|
|
visualXY: new Float32Array(bytes('visual-xy.binz').buffer),
|
|
};
|
|
if (data.indptr.length !== meta.neurons + 1 || data.targets.length !== meta.edges || data.weights.length !== meta.edges) {
|
|
throw new Error('FlyWire artifact lengths do not match metadata');
|
|
}
|
|
return data;
|
|
}
|
|
|
|
const realBase = join(fileURLToPath(new URL('../../../', import.meta.url)), 'data', 'fafb-v783');
|
|
const realDatasetPresent = existsSync(join(realBase, 'meta.json'));
|
|
|
|
test('LifNetwork is bit-exact with the legacy FlyBrain on the real FAFB v783 dataset', {
|
|
skip: realDatasetPresent ? false : 'data/fafb-v783/meta.json is absent in this worktree (dataset branch not merged)',
|
|
}, () => {
|
|
const data = loadRealDataset(realBase);
|
|
const legacy = new LegacyFlyBrain(data);
|
|
const next = new LifNetwork(data);
|
|
assert.equal(next.plasticity.edges.length, 16_384);
|
|
assert.deepEqual(next.plasticity.edges, legacy.plasticity.edges);
|
|
assert.equal(next.plasticity.exportState().topology, legacy.plasticity.exportState().topology);
|
|
assert.equal(next.plasticity.exportState().version, legacy.plasticity.exportState().version);
|
|
// The prototype predates the macro populations (`docs/design/macros.md` section 11), so it
|
|
// tracks fourteen roles where this network tracks forty-five. Everything else about the two is
|
|
// still asserted identical, which is the point of the test: the extra roles are counters over
|
|
// neurons that were always in the dataset and they change nothing that spikes.
|
|
const isMacro = (name: string) => name.startsWith('macro_');
|
|
assert.deepEqual(next.roleNames.filter(name => !isMacro(name)), legacy.roleNames);
|
|
// Thirty-one since sections 13 and 14; twenty-two when section 11 shipped.
|
|
assert.equal(next.roleNames.filter(isMacro).length, 31);
|
|
|
|
const frame = randomFrame(160, 144, 99);
|
|
legacy.setVisualFrame(frame); next.setVisualFrame(frame);
|
|
assert.equal(legacy.step(100), next.step(100));
|
|
legacy.reward(120); next.stimulate(120);
|
|
assert.equal(legacy.step(100), next.step(100));
|
|
legacy.plasticity.reinforce(1, legacy.ms); next.plasticity.reinforce(1, next.ms);
|
|
|
|
const withoutMacros = (rates: Record<string, number>) =>
|
|
Object.fromEntries(Object.entries(rates).filter(([name]) => !isMacro(name)));
|
|
assert.deepEqual({ ...next.exportState(), rates: withoutMacros(next.exportState().rates) }, legacy.exportState());
|
|
assert.deepEqual(withoutMacros(next.rates), legacy.rates);
|
|
// And the new roles are live rather than zero-filled padding.
|
|
assert.ok(next.roleNames.filter(isMacro).some(name => next.rates[name]! > 0));
|
|
assert.equal(next.populationRate, legacy.populationRate);
|
|
assert.equal(next.ms, 200);
|
|
assert.ok(next.populationRate > 0);
|
|
});
|