gm2d_workbench.js

import * as THREE from 'three';
import { ArcballControls } from 'three/addons/controls/ArcballControls.js';

const host = document.getElementById('gm2dThree');

if (host) {
    const ui = {
        train: document.getElementById('gm2dTrainMode'),
        sample: document.getElementById('gm2dSampleMode'),
        ode: document.getElementById('gm2dOdeMode'),
        sde: document.getElementById('gm2dSdeMode'),
        play: document.getElementById('gm2dPlay'),
        reset: document.getElementById('gm2dReset'),
        generate: document.getElementById('gm2dGenerate'),
        clear: document.getElementById('gm2dClear'),
        resetView: document.getElementById('gm2dResetView'),
        speed: document.getElementById('gm2dSpeed'),
        speedVal: document.getElementById('gm2dSpeedVal'),
        speedWrap: document.getElementById('gm2dSpeedWrap'),
        sigma: document.getElementById('gm2dSigma'),
        sigmaVal: document.getElementById('gm2dSigmaVal'),
        sigmaWrap: document.getElementById('gm2dSigmaWrap'),
        field: document.getElementById('gm2dShowField'),
        paths: document.getElementById('gm2dShowPaths'),
        x0: document.getElementById('gm2dShowX0'),
        x1: document.getElementById('gm2dShowX1'),
        stats: document.getElementById('gm2dStats')
    };

    const TIME_LENGTH = 8;
    const SPACE_EXTENT = 2.5;
    const SOURCE_SIGMA = 0.62;
    const PATH_STEPS = 100;
    const FIELD_REFRESH_EVERY = 50000;
    const MAX_ACCUMULATED = 5000;
    const MAX_ACTIVE = 72;
    const MAX_SDE_PATH_LINES = 50;
    const FIELD_COLOR = 0xeb5665;

    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 point(x, y) { return { x, y }; }
    function add(a, b) { return point(a.x + b.x, a.y + b.y); }
    function lerp(a, b, u) { return point(a.x + (b.x - a.x) * u, a.y + (b.y - a.y) * u); }
    function sourceSample() { return point(SOURCE_SIGMA * randn(), SOURCE_SIGMA * randn()); }

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

    function worldPoint(t, p) {
        return new THREE.Vector3((t - 0.5) * TIME_LENGTH, p.x, p.y);
    }

    const staticX0 = Array.from({ length: 150 }, sourceSample);
    const staticX1 = Array.from({ length: 220 }, (_, i) => targetSample((i + 0.35) / 220));
    const targetMean = staticX1.reduce((a, p) => add(a, p), point(0, 0));
    targetMean.x /= staticX1.length;
    targetMean.y /= staticX1.length;

    const model = {
        samples: new Float32Array(0),
        iterations: 0,
        emaLoss: null,
        lossHistory: [],

        reset() {
            this.samples = new Float32Array(0);
            this.iterations = 0;
            this.emaLoss = null;
            this.lossHistory = [];
        },

        applySnapshot(snapshot) {
            this.samples = snapshot.samples;
            this.iterations = snapshot.iterations;
            this.emaLoss = snapshot.emaLoss;
            this.lossHistory = snapshot.lossHistory;
        },

        targetBandwidth() {
            const fill = Math.min(1, this.samples.length / 2 / 1200);
            return 0.13 - 0.085 * fill;
        },

        /**
         * Conditional-mean bridge drift for an empirical target distribution.
         * Training supplies target examples; the Gaussian source and bridge give
         * p(x_t | x_1) in closed form, so the kernel only has to learn X_1.
         */
        predictU(t, p, process, sigma) {
            const omt = Math.max(1 - t, 0.005);
            if (!this.samples.length) {
                return point((targetMean.x - p.x) / omt, (targetMean.y - p.y) / omt);
            }
            const bridgeVariance = process === 'sde' ? sigma * sigma * t * (1 - t) : 0;
            const h = this.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 < this.samples.length; i += 2) {
                const dx = p.x - t * this.samples[i];
                const dy = p.y - t * this.samples[i + 1];
                minQ = Math.min(minQ, (dx * dx + dy * dy) / variance);
            }
            let den = 0, x1x = 0, x1y = 0;
            for (let i = 0; i < this.samples.length; i += 2) {
                const sx = this.samples[i], sy = this.samples[i + 1];
                const dx = p.x - t * sx;
                const 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;
            }
            const meanX1 = den > 0 ? point(x1x / den, x1y / den) : targetMean;
            return point((meanX1.x - p.x) / omt, (meanX1.y - p.y) / omt);
        }
    };

    const state = {
        mode: 'train',
        process: 'ode',
        running: false,
        speed: Number(ui.speed.value),
        sigma: Number(ui.sigma.value),
        lastFrame: performance.now(),
        lastFieldIteration: -FIELD_REFRESH_EVERY,
        active: [],
        arrived: [],
        autoSpawn: true,
        spawnTimer: 0,
        lastLossIteration: -1,
        slideActive: false,
        rendererSize: { w: 0, h: 0 }
    };

    let trainingWorker = null;
    let workerGeneration = 0;

    function ensureTrainingWorker() {
        if (trainingWorker) return trainingWorker;
        trainingWorker = new Worker(new URL('./gm2d_training_worker.js', import.meta.url), { type: 'module' });
        trainingWorker.addEventListener('message', ({ data }) => {
            if (data.type !== 'snapshot' || data.generation !== workerGeneration) return;
            model.applySnapshot(data);
            host.dataset.workerIterations = String(model.iterations);
        });
        return trainingWorker;
    }

    function trainingShouldRun() {
        return state.slideActive && state.mode === 'train' && state.running;
    }

    function syncTrainingWorker() {
        if (!trainingWorker) {
            if (trainingShouldRun()) resetTrainingWorker();
            return;
        }
        trainingWorker.postMessage({ type: 'run', running: trainingShouldRun() });
    }

    function resetTrainingWorker() {
        model.reset();
        workerGeneration++;
        ensureTrainingWorker().postMessage({
            type: 'reset',
            generation: workerGeneration,
            process: state.process,
            sigma: state.sigma,
            running: trainingShouldRun()
        });
    }

    host.replaceChildren();
    const scene = new THREE.Scene();
    scene.background = new THREE.Color(0xfbfdff);
    scene.fog = new THREE.Fog(0xfbfdff, 17, 30);

    const camera = new THREE.PerspectiveCamera(38, 1, 0.05, 80);
    camera.position.set(8.3, 5.5, 8.0);
    camera.lookAt(0, 0, 0);

    const renderer = new THREE.WebGLRenderer({ antialias: true, alpha: false });
    renderer.setPixelRatio(Math.min(window.devicePixelRatio || 1, 1.5));
    renderer.outputColorSpace = THREE.SRGBColorSpace;
    host.appendChild(renderer.domElement);
    const lossCanvas = document.createElement('canvas');
    lossCanvas.className = 'gm2d-loss';
    lossCanvas.width = 440;
    lossCanvas.height = 152;
    host.appendChild(lossCanvas);
    const lossCtx = lossCanvas.getContext('2d');
    const help = document.createElement('div');
    help.className = 'gm2d-help';
    help.textContent = 'drag rotate · wheel zoom · right-drag pan';
    host.appendChild(help);

    const controls = new ArcballControls(camera, renderer.domElement, scene);
    controls.enableAnimations = true;
    controls.enableGrid = false;
    controls.enableGizmos = false;
    controls.cursorZoom = true;
    controls.adjustNearFar = true;
    controls.minDistance = 5;
    controls.maxDistance = 28;
    controls.rotateSpeed = 1.05;
    controls.saveState();
    controls.setGizmosVisible(false);

    scene.add(new THREE.HemisphereLight(0xffffff, 0xc8d2df, 2.4));
    const keyLight = new THREE.DirectionalLight(0xffffff, 2.2);
    keyLight.position.set(3, 8, 9);
    scene.add(keyLight);

    const baseGroup = new THREE.Group();
    const fieldGroup = new THREE.Group();
    const pathGroup = new THREE.Group();
    const movingGroup = new THREE.Group();
    const arrivedGroup = new THREE.Group();
    scene.add(baseGroup, fieldGroup, pathGroup, movingGroup, arrivedGroup);
    const fieldArrows = [];
    const pathLines = [];

    function lineObject(points, color, opacity = 1) {
        const geometry = new THREE.BufferGeometry().setFromPoints(points);
        const material = new THREE.LineBasicMaterial({ color, transparent: opacity < 1, opacity });
        return new THREE.Line(geometry, material);
    }

    function labelSprite(text, color, scaleX = 1.6) {
        const canvas = document.createElement('canvas');
        canvas.width = 512;
        canvas.height = 96;
        const ctx = canvas.getContext('2d');
        ctx.clearRect(0, 0, canvas.width, canvas.height);
        ctx.font = '700 34px system-ui, sans-serif';
        ctx.textAlign = 'center';
        ctx.textBaseline = 'middle';
        ctx.fillStyle = color;
        ctx.fillText(text, 256, 48);
        const texture = new THREE.CanvasTexture(canvas);
        texture.colorSpace = THREE.SRGBColorSpace;
        const sprite = new THREE.Sprite(new THREE.SpriteMaterial({ map: texture, transparent: true, depthTest: false }));
        sprite.scale.set(scaleX, scaleX * 0.1875, 1);
        sprite.renderOrder = 8;
        return sprite;
    }

    function pointCloud(points, t, color, size, opacity) {
        const positions = new Float32Array(points.length * 3);
        points.forEach((p, i) => {
            const q = worldPoint(t, p);
            positions[i * 3] = q.x;
            positions[i * 3 + 1] = q.y;
            positions[i * 3 + 2] = q.z;
        });
        const geometry = new THREE.BufferGeometry();
        geometry.setAttribute('position', new THREE.BufferAttribute(positions, 3));
        const material = new THREE.PointsMaterial({ color, size, transparent: true, opacity, sizeAttenuation: true, depthWrite: false });
        return new THREE.Points(geometry, material);
    }

    function buildBase() {
        const timeStart = worldPoint(0, point(0, 0));
        const timeDir = new THREE.Vector3(1, 0, 0);
        baseGroup.add(new THREE.ArrowHelper(timeDir, timeStart, TIME_LENGTH + 0.15, 0x758392, 0.22, 0.12));

        for (const t of [0, 0.25, 0.5, 0.75, 1]) {
            const ring = [];
            for (let i = 0; i < 96; i++) {
                const a = i / 96 * Math.PI * 2;
                ring.push(worldPoint(t, point(SPACE_EXTENT * Math.cos(a), SPACE_EXTENT * Math.sin(a))));
            }
            ring.push(ring[0]);
            baseGroup.add(lineObject(ring, t === 0 ? 0x7db0df : t === 1 ? 0xe686b5 : 0xd7dee7, t === 0 || t === 1 ? 0.75 : 0.55));
            if (t > 0 && t < 1) {
                const l = labelSprite(t.toFixed(2), '#7a8794', 0.65);
                l.position.copy(worldPoint(t, point(-2.62, 0)));
                baseGroup.add(l);
            }
        }

        for (const t of [0, 1]) {
            baseGroup.add(lineObject([worldPoint(t, point(-SPACE_EXTENT, 0)), worldPoint(t, point(SPACE_EXTENT, 0))], 0xdfe5ec, 0.8));
            baseGroup.add(lineObject([worldPoint(t, point(0, -SPACE_EXTENT)), worldPoint(t, point(0, SPACE_EXTENT))], 0xdfe5ec, 0.8));
        }

        const sourceDisc = new THREE.Mesh(
            new THREE.CircleGeometry(SPACE_EXTENT, 64),
            new THREE.MeshBasicMaterial({ color: 0xddebf8, transparent: true, opacity: 0.23, side: THREE.DoubleSide, depthWrite: false })
        );
        sourceDisc.rotation.y = Math.PI / 2;
        sourceDisc.position.x = -TIME_LENGTH / 2;
        const targetDisc = new THREE.Mesh(
            new THREE.CircleGeometry(SPACE_EXTENT, 64),
            new THREE.MeshBasicMaterial({ color: 0xf9dce9, transparent: true, opacity: 0.2, side: THREE.DoubleSide, depthWrite: false })
        );
        targetDisc.rotation.y = Math.PI / 2;
        targetDisc.position.x = TIME_LENGTH / 2;
        baseGroup.add(sourceDisc, targetDisc);

        const timeLabel = labelSprite('time  t', '#667583', 1.25);
        timeLabel.position.set(0, -0.28, -0.15);
        const x0Label = labelSprite('X₀  Gaussian', '#0056b3', 1.55);
        x0Label.position.set(-TIME_LENGTH / 2, -2.82, 0);
        const x1Label = labelSprite('X₁  spiral', '#d63384', 1.45);
        x1Label.position.set(TIME_LENGTH / 2, 2.82, 0);
        baseGroup.add(timeLabel, x0Label, x1Label);
    }

    const sourcePoints = pointCloud(staticX0, 0, 0x2677c9, 0.075, 0.62);
    const targetPoints = pointCloud(staticX1, 1, 0xd63384, 0.075, 0.7);
    scene.add(sourcePoints, targetPoints);
    buildBase();

    function disposeGroup(group) {
        while (group.children.length) {
            const obj = group.children.pop();
            obj.geometry?.dispose();
            if (Array.isArray(obj.material)) obj.material.forEach(m => m.dispose());
            else obj.material?.dispose();
        }
    }

    function updateField() {
        const values = [-1.8, -0.6, 0.6, 1.8];
        let index = 0;
        for (const t of [0.05, 0.16, 0.28, 0.40, 0.52, 0.64, 0.76, 0.86, 0.94]) {
            for (const x of values) for (const y of values) {
                const p = point(x, y);
                const u = model.predictU(t, p, state.process, state.sigma);
                const tangent = new THREE.Vector3(TIME_LENGTH, u.x, u.y).normalize();
                let arrow = fieldArrows[index++];
                if (!arrow) {
                    arrow = new THREE.ArrowHelper(tangent, worldPoint(t, p), 0.38, FIELD_COLOR, 0.105, 0.06);
                    fieldArrows.push(arrow);
                    fieldGroup.add(arrow);
                } else {
                    arrow.position.copy(worldPoint(t, p));
                    arrow.setDirection(tangent);
                    arrow.setLength(0.38, 0.105, 0.06);
                }
            }
        }
        state.lastFieldIteration = model.iterations;
    }

    function probabilitySeeds(n) {
        const out = [];
        const golden = Math.PI * (3 - Math.sqrt(5));
        for (let i = 0; i < n; i++) {
            const q = Math.min(0.985, (i + 0.5) / n);
            const r = SOURCE_SIGMA * Math.sqrt(-2 * Math.log(1 - q));
            out.push(point(r * Math.cos(i * golden), r * Math.sin(i * golden)));
        }
        return out;
    }
    const odeSeeds = probabilitySeeds(40);

    function integrate(seed) {
        const pts = [{ t: 0, p: point(seed.x, seed.y) }];
        let p = point(seed.x, seed.y);
        const dt = 1 / PATH_STEPS;
        for (let i = 0; i < PATH_STEPS; i++) {
            const t = i / PATH_STEPS;
            const t1 = (i + 1) / PATH_STEPS;
            if (state.process === 'ode') {
                // Midpoint integration keeps the learned spiral flow smooth near t=1.
                const u0 = model.predictU(t, p, state.process, state.sigma);
                const midpoint = point(p.x + 0.5 * dt * u0.x, p.y + 0.5 * dt * u0.y);
                const um = model.predictU(t + 0.5 * dt, midpoint, state.process, state.sigma);
                p = point(p.x + dt * um.x, p.y + dt * um.y);
            } else {
                const u = model.predictU(t, p, state.process, state.sigma);
                p = point(
                    p.x + u.x * dt + state.sigma * Math.sqrt(dt) * randn(),
                    p.y + u.y * dt + state.sigma * Math.sqrt(dt) * randn()
                );
            }
            pts.push({ t: t1, p });
        }
        return pts;
    }

    const dotGeometry = new THREE.SphereGeometry(0.075, 14, 10);

    function buildOdeReferenceLines() {
        for (const seed of odeSeeds) {
            const pts = integrate(seed);
            const path = lineObject(pts.map(v => worldPoint(v.t, v.p)), 0x7384cf, 0.28);
            pathGroup.add(path);
            pathLines.push(path);
        }
    }

    function addSamples(n) {
        const remaining = MAX_ACCUMULATED - state.arrived.length - state.active.length;
        n = Math.max(0, Math.min(n, MAX_ACTIVE - state.active.length, remaining));
        const color = state.process === 'ode' ? 0x8656d6 : 0x9a6740;
        for (let i = 0; i < n; i++) {
            // Animated samples are always genuine random draws. The regular ODE
            // seeds above are used only to make stable reference streamlines.
            const seed = sourceSample();
            const pts = integrate(seed);
            let path = null;
            if (state.process === 'sde') {
                path = lineObject(pts.map(v => worldPoint(v.t, v.p)), 0xaa7b55, 0.42);
                pathGroup.add(path);
                pathLines.push(path);
                while (pathLines.length > MAX_SDE_PATH_LINES) {
                    const old = pathLines.shift();
                    pathGroup.remove(old);
                    old.geometry.dispose();
                    old.material.dispose();
                }
            }
            const dot = new THREE.Mesh(dotGeometry, new THREE.MeshPhongMaterial({ color, shininess: 70 }));
            dot.position.copy(worldPoint(0, seed));
            movingGroup.add(dot);
            state.active.push({ pts, progress: 0, dot, path });
        }
    }

    function rebuildArrived() {
        disposeGroup(arrivedGroup);
        if (!state.arrived.length) return;
        arrivedGroup.add(pointCloud(state.arrived, 1, state.process === 'ode' ? 0x3d9b6f : 0x8c603e, 0.105, 0.92));
    }

    function clearSamples() {
        for (const item of state.active) item.dot.material.dispose();
        state.active = [];
        state.arrived = [];
        disposeGroup(pathGroup);
        pathLines.length = 0;
        movingGroup.clear();
        disposeGroup(arrivedGroup);
        state.spawnTimer = 0;
    }

    function enterSampling() {
        state.mode = 'sample';
        state.running = true;
        syncTrainingWorker();
        clearSamples();
        state.autoSpawn = true;
        if (state.process === 'ode') buildOdeReferenceLines();
        addSamples(state.process === 'ode' ? 32 : 28);
        syncUi();
    }

    function enterTraining() {
        state.mode = 'train';
        state.running = false;
        clearSamples();
        syncUi();
        syncTrainingWorker();
    }

    function setProcess(process) {
        if (state.process === process) return;
        state.process = process;
        state.lastLossIteration = -1;
        state.lastFieldIteration = -FIELD_REFRESH_EVERY;
        enterTraining();
        resetTrainingWorker();
        if (ui.field.checked) updateField();
    }

    function resetModel() {
        state.running = false;
        state.lastLossIteration = -1;
        state.lastFieldIteration = -FIELD_REFRESH_EVERY;
        clearSamples();
        resetTrainingWorker();
        if (ui.field.checked) updateField();
        syncUi();
    }

    function syncUi() {
        ui.train.classList.toggle('active', state.mode === 'train');
        ui.sample.classList.toggle('active', state.mode === 'sample');
        ui.ode.classList.toggle('active', state.process === 'ode');
        ui.sde.classList.toggle('active', state.process === 'sde');
        ui.play.textContent = state.mode === 'train'
            ? (state.running ? 'Pause' : 'Train')
            : (state.running ? '⏸' : '▶');
        ui.play.title = state.mode === 'train'
            ? (state.running ? 'Pause training' : 'Start training')
            : (state.running ? 'Pause sampling' : 'Resume sampling');
        ui.speedVal.textContent = `${state.speed}×`;
        ui.sigmaVal.textContent = state.sigma.toFixed(2);
        ui.sigma.disabled = state.process !== 'sde';
        ui.sigmaWrap.style.opacity = state.process === 'sde' ? '1' : '0.42';
        ui.generate.hidden = state.mode !== 'sample';
        ui.clear.hidden = state.mode !== 'sample';
        ui.speedWrap.hidden = state.mode !== 'sample';
        sourcePoints.visible = ui.x0.checked;
        targetPoints.visible = ui.x1.checked;
        fieldGroup.visible = ui.field.checked;
        pathGroup.visible = state.mode === 'sample' && ui.paths.checked;
        movingGroup.visible = state.mode === 'sample';
        arrivedGroup.visible = state.mode === 'sample';
        lossCanvas.hidden = state.mode !== 'train';
    }

    function drawLoss() {
        const history = model.lossHistory;
        const newestIteration = history.length ? history[history.length - 1].iteration : model.iterations;
        if (state.lastLossIteration === newestIteration) return;
        state.lastLossIteration = newestIteration;
        const w = lossCanvas.width, h = lossCanvas.height;
        lossCtx.clearRect(0, 0, w, h);
        lossCtx.fillStyle = 'rgba(255,255,255,0.92)';
        lossCtx.fillRect(0, 0, w, h);
        lossCtx.font = '600 20px system-ui, sans-serif';
        lossCtx.fillStyle = '#5f6e7d';
        lossCtx.fillText(`${state.process.toUpperCase()} training loss`, 18, 28);
        const left = 18, right = w - 14, top = 40, bottom = h - 18;
        lossCtx.strokeStyle = '#d9e0e8';
        lossCtx.lineWidth = 2;
        lossCtx.beginPath();
        lossCtx.moveTo(left, bottom);
        lossCtx.lineTo(right, bottom);
        lossCtx.stroke();
        if (history.length < 2) return;
        const maxLoss = Math.max(0.25, ...history.map(v => v.loss));
        lossCtx.strokeStyle = state.process === 'ode' ? '#d94d77' : '#986642';
        lossCtx.lineWidth = 3;
        lossCtx.beginPath();
        history.forEach((v, i) => {
            const x = left + i / Math.max(1, history.length - 1) * (right - left);
            const y = bottom - Math.min(1, v.loss / maxLoss) * (bottom - top);
            if (i === 0) lossCtx.moveTo(x, y);
            else lossCtx.lineTo(x, y);
        });
        lossCtx.stroke();
    }

    function updateStats() {
        host.dataset.iterations = String(model.iterations);
        host.dataset.fieldIteration = String(state.lastFieldIteration);
        host.dataset.activeSamples = String(state.active.length);
        host.dataset.arrivedSamples = String(state.arrived.length);
        host.dataset.pathLines = String(pathLines.length);
        host.dataset.fieldArrows = String(fieldArrows.length);
        if (state.mode === 'train') {
            const loss = model.emaLoss == null ? '—' : model.emaLoss.toFixed(3);
            const until = FIELD_REFRESH_EVERY - (model.iterations - state.lastFieldIteration);
            const fieldStatus = ui.field.checked ? `field refresh ${Math.max(0, until)}` : 'field updates off';
            ui.stats.textContent = `training model · iter ${model.iterations} · loss ${loss} · ${fieldStatus}`;
        } else {
            ui.stats.textContent = `${state.active.length} moving · ${state.arrived.length} accumulated · trained ${model.iterations}`;
        }
    }

    function resizeRenderer() {
        const w = Math.max(1, host.clientWidth);
        const h = Math.max(1, host.clientHeight);
        if (w === state.rendererSize.w && h === state.rendererSize.h) return;
        state.rendererSize = { w, h };
        renderer.setSize(w, h, false);
        camera.aspect = w / h;
        camera.updateProjectionMatrix();
        controls.setCamera(camera);
    }

    function tick(now) {
        // A resume click can occur just after the current rAF timestamp was
        // captured, so clamp the one possible negative frame to zero.
        const dt = Math.max(0, Math.min(0.05, (now - state.lastFrame) / 1000));
        state.lastFrame = now;
        resizeRenderer();
        if (state.running && state.mode === 'train') {
            if (ui.field.checked && model.iterations - state.lastFieldIteration >= FIELD_REFRESH_EVERY) updateField();
        } else if (state.running && state.mode === 'sample') {
            const rate = 0.18 + state.speed / 58;
            for (let i = state.active.length - 1; i >= 0; i--) {
                const item = state.active[i];
                item.progress = Math.min(1, item.progress + dt * rate);
                if (item.progress >= 1) {
                    movingGroup.remove(item.dot);
                    item.dot.material.dispose();
                    state.arrived.push(item.pts[item.pts.length - 1].p);
                    state.active.splice(i, 1);
                    continue;
                }
                const f = item.progress * (item.pts.length - 1);
                const j = Math.min(item.pts.length - 2, Math.floor(f));
                const u = f - j;
                const p = lerp(item.pts[j].p, item.pts[j + 1].p, u);
                const t = item.pts[j].t + (item.pts[j + 1].t - item.pts[j].t) * u;
                item.dot.position.copy(worldPoint(t, p));
            }
            if (state.autoSpawn && (state.active.length === 0 || state.spawnTimer <= 0)) {
                rebuildArrived();
                if (state.arrived.length + state.active.length < MAX_ACCUMULATED) addSamples(8);
                state.spawnTimer = 1.7;
            }
            state.spawnTimer -= dt;
        }
        drawLoss();
        updateStats();
        renderer.render(scene, camera);
    }

    ui.train.addEventListener('click', enterTraining);
    ui.sample.addEventListener('click', enterSampling);
    ui.ode.addEventListener('click', () => setProcess('ode'));
    ui.sde.addEventListener('click', () => setProcess('sde'));
    ui.play.addEventListener('click', () => {
        state.running = !state.running;
        state.lastFrame = performance.now();
        syncTrainingWorker();
        syncUi();
    });
    ui.reset.addEventListener('click', resetModel);
    ui.generate.addEventListener('click', () => { state.autoSpawn = true; addSamples(16); });
    ui.clear.addEventListener('click', () => { clearSamples(); state.autoSpawn = false; });
    ui.resetView.addEventListener('click', () => controls.reset());
    ui.speed.addEventListener('input', () => { state.speed = Number(ui.speed.value); syncUi(); });
    ui.sigma.addEventListener('input', () => { state.sigma = Number(ui.sigma.value); syncUi(); });
    ui.sigma.addEventListener('change', () => { if (state.process === 'sde') resetModel(); });
    ui.field.addEventListener('change', () => {
        if (ui.field.checked && model.iterations !== state.lastFieldIteration) updateField();
        syncUi();
    });
    ui.paths.addEventListener('change', syncUi);
    ui.x0.addEventListener('change', syncUi);
    ui.x1.addEventListener('change', syncUi);

    if (ui.field.checked) updateField();
    syncUi();
    drawLoss();
    resizeRenderer();
    renderer.render(scene, camera);

    const renderLoop = window.makeSlideRafLoop(tick, {
        onStart: () => { state.lastFrame = performance.now(); resizeRenderer(); }
    });
    const workbenchSlide = window.slideIndexOf(host);
    window.SlideAnim.register(workbenchSlide, renderLoop);
    window.SlideAnim.register([workbenchSlide, workbenchSlide + 1, workbenchSlide + 2], {
        start() {
            state.slideActive = true;
            syncTrainingWorker();
        },
        stop() {
            state.slideActive = false;
            syncTrainingWorker();
        }
    });
    window.SlideAnim.sync(window.slideIndexOf(document.querySelector('.slide.active')));
}