sphere_flow_training_worker.js

const MAX_SAMPLES = 2400;
const STRIDE = 7;
const SNAPSHOT_INTERVAL_MS = 160;
const WORK_CHUNK_MS = 13;
const SOURCE_CENTER = [-1, 0, 0];
const TARGET_CENTER = [1, 0, 0];

let samples = new Float32Array(MAX_SAMPLES * STRIDE);
let sampleCount = 0;
let sampleCursor = 0;
let iterations = 0;
let emaLoss = null;
let generation = 0;
let running = false;
let runToken = 0;
let lastSnapshotAt = 0;

const dot = (a, b) => a[0] * b[0] + a[1] * b[1] + a[2] * b[2];
const norm = a => Math.hypot(a[0], a[1], a[2]);
const scale = (a, s) => [a[0] * s, a[1] * s, a[2] * s];
const add = (a, b) => [a[0] + b[0], a[1] + b[1], a[2] + b[2]];
const normalize = a => { const n = Math.max(1e-12, norm(a)); return scale(a, 1 / n); };

function randn() {
    let u = 0, v = 0;
    while (u === 0) u = Math.random();
    while (v === 0) v = Math.random();
    return Math.sqrt(-2 * Math.log(u)) * Math.cos(2 * Math.PI * v);
}

function tangentProject(x, v) { return add(v, scale(x, -dot(x, v))); }

function expMap(x, v) {
    const speed = norm(v);
    if (speed < 1e-10) return [...x];
    return normalize(add(scale(x, Math.cos(speed)), scale(v, Math.sin(speed) / speed)));
}

function logMap(x, y) {
    const cosine = Math.max(-0.999999, Math.min(0.999999, dot(x, y)));
    const theta = Math.acos(cosine);
    if (theta < 1e-8) return [0, 0, 0];
    return scale(add(y, scale(x, -cosine)), theta / Math.max(1e-8, Math.sin(theta)));
}

function parallelTransport(p, q, v) {
    const denom = Math.max(1e-5, 1 + dot(p, q));
    return tangentProject(q, add(v, scale(add(p, q), -dot(v, q) / denom)));
}

function sourceSample() {
    const raw = [0, 0.42 * randn(), 0.42 * randn()];
    return expMap(SOURCE_CENTER, raw);
}

function targetPoint(u, jitter = 0) {
    const radius = 0.16 + 0.78 * u + jitter;
    const angle = -0.55 + 5.2 * Math.PI * u;
    return expMap(TARGET_CENTER, [0, radius * Math.cos(angle), radius * Math.sin(angle)]);
}

function coupledTarget(x0) {
    // The angle around the isotropic source cap is uniform, so this is a
    // measure-preserving coupling to the target's uniform spiral parameter.
    const sourceTangent = logMap(SOURCE_CENTER, x0);
    const u = (Math.atan2(sourceTangent[2], sourceTangent[1]) + Math.PI) / (2 * Math.PI);
    const radialJitter = 0.07 * Math.max(-1.5, Math.min(1.5, norm(sourceTangent) / 0.42 - 1.2));
    return targetPoint(u, radialJitter);
}

function bandwidths() {
    const fill = Math.min(1, sampleCount / 2000);
    return { time: 0.14 - 0.10 * fill, angle: 0.55 - 0.45 * fill };
}

function predict(t, x) {
    if (!sampleCount) return scale(logMap(x, [0, 1, 0]), 1 / Math.max(1 - t, 0.01));
    const bw = bandwidths();
    let den = 0, out = [0, 0, 0];
    for (let i = 0; i < sampleCount; i++) {
        const k = i * STRIDE;
        const p = [samples[k + 1], samples[k + 2], samples[k + 3]];
        const dt = (t - samples[k]) / bw.time;
        const angle = Math.acos(Math.max(-1, Math.min(1, dot(x, p)))) / bw.angle;
        const q = dt * dt + angle * angle;
        if (q > 18) continue;
        const w = Math.exp(-0.5 * q);
        const u = parallelTransport(p, x, [samples[k + 4], samples[k + 5], samples[k + 6]]);
        out = add(out, scale(u, w)); den += w;
    }
    return den > 1e-8 ? scale(tangentProject(x, out), 1 / den) : [0, 0, 0];
}

function trainOne() {
    const x0 = sourceSample(), x1 = coupledTarget(x0);
    const t = 0.005 + 0.99 * Math.random();
    const xt = expMap(x0, scale(logMap(x0, x1), t));
    const target = scale(logMap(xt, x1), 1 / Math.max(1 - t, 0.01));
    const predicted = predict(t, xt);
    const loss = (target[0] - predicted[0]) ** 2 + (target[1] - predicted[1]) ** 2 + (target[2] - predicted[2]) ** 2;
    const k = sampleCursor * STRIDE;
    samples[k] = t;
    samples[k + 1] = xt[0]; samples[k + 2] = xt[1]; samples[k + 3] = xt[2];
    samples[k + 4] = target[0]; samples[k + 5] = target[1]; samples[k + 6] = target[2];
    sampleCursor = (sampleCursor + 1) % MAX_SAMPLES;
    sampleCount = Math.min(MAX_SAMPLES, sampleCount + 1);
    iterations++;
    emaLoss = emaLoss == null ? loss : 0.985 * emaLoss + 0.015 * loss;
}

function postSnapshot() {
    const packed = samples.slice(0, sampleCount * STRIDE);
    postMessage({ type: 'snapshot', generation, iterations, emaLoss, samples: packed }, [packed.buffer]);
    lastSnapshotAt = performance.now();
}

function pump(token) {
    if (!running || token !== runToken) return;
    const deadline = performance.now() + WORK_CHUNK_MS;
    do trainOne(); while (performance.now() < deadline);
    if (performance.now() - lastSnapshotAt >= SNAPSHOT_INTERVAL_MS) postSnapshot();
    setTimeout(() => pump(token), 0);
}

function setRunning(next) {
    running = next; runToken++;
    if (running) pump(runToken); else postSnapshot();
}

self.onmessage = ({ data }) => {
    if (data.type === 'reset') {
        generation = data.generation; samples = new Float32Array(MAX_SAMPLES * STRIDE);
        sampleCount = 0; sampleCursor = 0; iterations = 0; emaLoss = null; lastSnapshotAt = 0;
        postSnapshot(); setRunning(data.running);
    } else if (data.type === 'run') setRunning(data.running);
};