gm2d_training_worker.js

const SOURCE_SIGMA = 0.62;
const MAX_SAMPLES = 1600;
const SNAPSHOT_INTERVAL_MS = 140;
const WORK_CHUNK_MS = 14;

let processMode = 'ode';
let sigma = 0.45;
let generation = 0;
let running = false;
let runToken = 0;
let iterations = 0;
let emaLoss = null;
let lossHistory = [];
let sampleCount = 0;
let sampleCursor = 0;
let samples = new Float32Array(MAX_SAMPLES * 2);
let lastSnapshotAt = 0;

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 sourceSample() {
    return { x: SOURCE_SIGMA * randn(), y: SOURCE_SIGMA * randn() };
}

function targetSample() {
    const u = Math.random();
    const theta = 0.75 * Math.PI + u * 3.35 * Math.PI;
    const r = 0.22 + 1.95 * u;
    const radial = 0.055 * randn();
    const tangent = 0.035 * randn();
    const rr = r + radial;
    return {
        x: rr * Math.cos(theta) - tangent * Math.sin(theta),
        y: rr * Math.sin(theta) + tangent * Math.cos(theta)
    };
}

function targetBandwidth() {
    const fill = Math.min(1, sampleCount / 1200);
    return 0.13 - 0.085 * fill;
}

function predictU(t, p) {
    const omt = Math.max(1 - t, 0.005);
    if (!sampleCount) return { x: -p.x / omt, y: -p.y / omt };
    const bridgeVariance = processMode === 'sde' ? sigma * sigma * t * (1 - t) : 0;
    const h = targetBandwidth();
    const variance = Math.max(1e-6,
        (1 - t) * (1 - t) * SOURCE_SIGMA * SOURCE_SIGMA + bridgeVariance + t * t * h * h
    );
    let minQ = Infinity;
    for (let i = 0; i < sampleCount; i++) {
        const dx = p.x - t * samples[2 * i];
        const dy = p.y - t * samples[2 * i + 1];
        minQ = Math.min(minQ, (dx * dx + dy * dy) / variance);
    }
    let den = 0, x1x = 0, x1y = 0;
    for (let i = 0; i < sampleCount; i++) {
        const sx = samples[2 * i], sy = samples[2 * i + 1];
        const dx = p.x - t * sx, dy = p.y - t * sy;
        const q = (dx * dx + dy * dy) / variance;
        if (q - minQ > 32) continue;
        const w = Math.exp(-0.5 * (q - minQ));
        den += w; x1x += w * sx; x1y += w * sy;
    }
    return { x: (x1x / den - p.x) / omt, y: (x1y / den - p.y) / omt };
}

function trainOne() {
    const x0 = sourceSample();
    const x1 = targetSample();
    const t = 0.01 + 0.95 * Math.random();
    const xt = { x: (1 - t) * x0.x + t * x1.x, y: (1 - t) * x0.y + t * x1.y };
    if (processMode === 'sde') {
        const sd = sigma * Math.sqrt(t * (1 - t));
        xt.x += sd * randn(); xt.y += sd * randn();
    }
    const u = predictU(t, xt);
    const denom = Math.max(1 - t, 0.005);
    const ux = (x1.x - xt.x) / denom, uy = (x1.y - xt.y) / denom;
    const loss = (u.x - ux) ** 2 + (u.y - uy) ** 2;
    samples[2 * sampleCursor] = x1.x;
    samples[2 * sampleCursor + 1] = x1.y;
    sampleCursor = (sampleCursor + 1) % MAX_SAMPLES;
    sampleCount = Math.min(MAX_SAMPLES, sampleCount + 1);
    iterations++;
    emaLoss = emaLoss == null ? loss : 0.985 * emaLoss + 0.015 * loss;
    if (iterations % 20 === 0) {
        lossHistory.push({ iteration: iterations, loss: emaLoss });
        if (lossHistory.length > 140) lossHistory.shift();
    }
}

function postSnapshot() {
    const packed = samples.slice(0, sampleCount * 2);
    postMessage({ type: 'snapshot', generation, iterations, emaLoss, lossHistory, 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;
        processMode = data.process;
        sigma = data.sigma;
        iterations = 0;
        emaLoss = null;
        lossHistory = [];
        sampleCount = 0;
        sampleCursor = 0;
        samples = new Float32Array(MAX_SAMPLES * 2);
        lastSnapshotAt = 0;
        postSnapshot();
        setRunning(data.running);
    } else if (data.type === 'run') {
        setRunning(data.running);
    }
};