sphere_flow_workbench.js

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

const host = document.getElementById('sphereFlowThree');
if (host) {
    const ui = {
        train: document.getElementById('sphereTrain'), reset: document.getElementById('sphereReset'),
        generate: document.getElementById('sphereGenerate'), clear: document.getElementById('sphereClear'),
        view: document.getElementById('sphereView'), time: document.getElementById('sphereTime'),
        timeVal: document.getElementById('sphereTimeVal'), field: document.getElementById('sphereShowField'),
        paths: document.getElementById('sphereShowPaths'), x0: document.getElementById('sphereShowX0'),
        x1: document.getElementById('sphereShowX1'), stats: document.getElementById('sphereStats')
    };
    const R = 2.25, STRIDE = 7, FIELD_REFRESH = 12000, PATH_STEPS = 100;
    const SOURCE_CENTER = [-1, 0, 0], TARGET_CENTER = [1, 0, 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);};
    const tangentProject=(x,v)=>add(v,scale(x,-dot(x,v)));
    function randn(){let u=0,v=0;while(!u)u=Math.random();while(!v)v=Math.random();return Math.sqrt(-2*Math.log(u))*Math.cos(2*Math.PI*v);}
    function expMap(x,v){const s=norm(v);if(s<1e-10)return [...x];return normalize(add(scale(x,Math.cos(s)),scale(v,Math.sin(s)/s)));}
    function logMap(x,y){const c=Math.max(-.999999,Math.min(.999999,dot(x,y))),a=Math.acos(c);return a<1e-8?[0,0,0]:scale(add(y,scale(x,-c)),a/Math.max(1e-8,Math.sin(a)));}
    function parallelTransport(p,q,v){const d=Math.max(1e-5,1+dot(p,q));return tangentProject(q,add(v,scale(add(p,q),-dot(v,q)/d)));}
    function sourceSample(){return expMap(SOURCE_CENTER,[0,.42*randn(),.42*randn()]);}
    function targetPoint(u,jitter=0){const radius=.16+.78*u+jitter,angle=-.55+5.2*Math.PI*u;return expMap(TARGET_CENTER,[0,radius*Math.cos(angle),radius*Math.sin(angle)]);}
    const world=x=>new THREE.Vector3(R*x[0],R*x[1],R*x[2]);

    const model={samples:new Float32Array(0),iterations:0,emaLoss:null,
        reset(){this.samples=new Float32Array(0);this.iterations=0;this.emaLoss=null;},
        bandwidths(){const f=Math.min(1,this.samples.length/STRIDE/2000);return{time:.14-.10*f,angle:.55-.45*f};},
        predict(t,x){
            if(!this.samples.length)return scale(logMap(x,[0,1,0]),1/Math.max(1-t,.01));
            const bw=this.bandwidths();let den=0,out=[0,0,0],bestQ=Infinity,best=null;
            for(let k=0;k<this.samples.length;k+=STRIDE){const p=[this.samples[k+1],this.samples[k+2],this.samples[k+3]];const dt=(t-this.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<bestQ){bestQ=q;best=k;}if(q>18)continue;const w=Math.exp(-.5*q);const u=parallelTransport(p,x,[this.samples[k+4],this.samples[k+5],this.samples[k+6]]);out=add(out,scale(u,w));den+=w;}
            if(den>1e-8)return scale(tangentProject(x,out),1/den);
            const p=[this.samples[best+1],this.samples[best+2],this.samples[best+3]];return parallelTransport(p,x,[this.samples[best+4],this.samples[best+5],this.samples[best+6]]);
        }
    };
    const state={running:false,scopeActive:false,lastFrame:performance.now(),lastField:-FIELD_REFRESH,active:[],arrived:[],paths:[],rendererSize:{w:0,h:0}};

    let worker=null,generation=0;
    function ensureWorker(){if(worker)return worker;worker=new Worker(new URL('./sphere_flow_training_worker.js',import.meta.url),{type:'module'});worker.onmessage=({data})=>{if(data.type!=='snapshot'||data.generation!==generation)return;model.samples=data.samples;model.iterations=data.iterations;model.emaLoss=data.emaLoss;host.dataset.workerIterations=String(data.iterations);ui.generate.disabled=model.samples.length<STRIDE*20;};return worker;}
    const shouldTrain=()=>state.scopeActive&&state.running;
    function syncWorker(){if(!worker){if(shouldTrain())resetWorker();return;}worker.postMessage({type:'run',running:shouldTrain()});}
    function resetWorker(){model.reset();generation++;ensureWorker().postMessage({type:'reset',generation,running:shouldTrain()});}

    host.replaceChildren();
    const scene=new THREE.Scene();scene.background=new THREE.Color(0xfbfdff);
    const camera=new THREE.PerspectiveCamera(38,1,.05,80);camera.position.set(5.4,3.5,5.2);camera.lookAt(0,0,0);
    const renderer=new THREE.WebGLRenderer({antialias:true});renderer.setPixelRatio(Math.min(devicePixelRatio||1,1.5));renderer.outputColorSpace=THREE.SRGBColorSpace;host.appendChild(renderer.domElement);
    scene.add(new THREE.HemisphereLight(0xffffff,0x657080,2.1));const light=new THREE.DirectionalLight(0xffffff,2.2);light.position.set(4,6,5);scene.add(light);
    const sphere=new THREE.Mesh(new THREE.SphereGeometry(R,48,32),new THREE.MeshPhongMaterial({color:0xdceafb,transparent:true,opacity:.33,side:THREE.DoubleSide,depthWrite:false,shininess:55}));scene.add(sphere);
    const wire=new THREE.LineSegments(new THREE.WireframeGeometry(new THREE.SphereGeometry(R*1.002,24,16)),new THREE.LineBasicMaterial({color:0x9fb5cb,transparent:true,opacity:.2}));scene.add(wire);
    const controls=new ArcballControls(camera,renderer.domElement,scene);controls.enableAnimations=true;controls.enableGizmos=false;controls.setGizmosVisible(false);controls.minDistance=4;controls.maxDistance=14;controls.saveState();
    const help=document.createElement('div');help.className='gm2d-help';help.textContent='drag rotate · wheel zoom · right-drag pan';host.appendChild(help);

    function pointsObject(points,color,size,opacity=.8){const g=new THREE.BufferGeometry().setFromPoints(points.map(world));const m=new THREE.PointsMaterial({color,size,transparent:true,opacity,depthWrite:false,sizeAttenuation:true});return new THREE.Points(g,m);}
    const sourcePoints=pointsObject(Array.from({length:180},sourceSample),0x2677c9,.085,.68);
    const targetData=Array.from({length:320},()=>targetPoint(Math.random(),.04*randn()));
    const targetCurve=Array.from({length:360},(_,i)=>targetPoint(i/359));
    const targetPoints=pointsObject(targetData,0xd63384,.075,.82);
    const targetLine=new THREE.Line(new THREE.BufferGeometry().setFromPoints(targetCurve.map(world)),new THREE.LineBasicMaterial({color:0xd63384,transparent:true,opacity:.72}));
    scene.add(sourcePoints,targetPoints,targetLine);
    const fieldGroup=new THREE.Group(),pathGroup=new THREE.Group(),movingGroup=new THREE.Group(),arrivedGroup=new THREE.Group();scene.add(fieldGroup,pathGroup,movingGroup,arrivedGroup);
    const fieldArrows=[];
    const fieldSites=Array.from({length:78},(_,i)=>{const y=1-2*(i+.5)/78,r=Math.sqrt(Math.max(0,1-y*y)),a=i*Math.PI*(3-Math.sqrt(5));return[r*Math.cos(a),y,r*Math.sin(a)];});
    function updateField(){const t=Number(ui.time.value);fieldSites.forEach((x,i)=>{const u=model.predict(t,x),n=norm(u);const tangent=n>1e-7?scale(u,1/n):normalize(tangentProject(x,[0,1,0]));let ar=fieldArrows[i];const origin=world(scale(x,1.012));const dir=new THREE.Vector3(...tangent);if(!ar){ar=new THREE.ArrowHelper(dir,origin,.32,0xeb5665,.09,.05);fieldArrows.push(ar);fieldGroup.add(ar);}else{ar.position.copy(origin);ar.setDirection(dir);ar.setLength(.32,.09,.05);}});state.lastField=model.iterations;}

    function integrate(seed){const pts=[seed];let x=[...seed],dt=1/PATH_STEPS;for(let i=0;i<PATH_STEPS;i++){const t=i/PATH_STEPS;const u0=model.predict(t,x);const mid=expMap(x,scale(u0,.5*dt));const um=model.predict(t+.5*dt,mid);x=expMap(x,scale(um,dt));pts.push(x);}return pts;}
    const dotGeometry=new THREE.SphereGeometry(.085,12,9);
    function lineObject(pts){return new THREE.Line(new THREE.BufferGeometry().setFromPoints(pts.map(world)),new THREE.LineBasicMaterial({color:0x596fc7,transparent:true,opacity:.42}));}
    function addSamples(n){if(model.samples.length<STRIDE*20)return;for(let i=0;i<n&&state.active.length<64;i++){const pts=integrate(sourceSample()),line=lineObject(pts),dotMesh=new THREE.Mesh(dotGeometry,new THREE.MeshPhongMaterial({color:0x8257d6,shininess:65}));dotMesh.position.copy(world(pts[0]));pathGroup.add(line);movingGroup.add(dotMesh);state.paths.push(line);while(state.paths.length>80){const old=state.paths.shift();pathGroup.remove(old);old.geometry.dispose();old.material.dispose();}state.active.push({pts,line,dot:dotMesh,progress:0});}}
    function rebuildArrived(){
        while(arrivedGroup.children.length){const o=arrivedGroup.children.pop();o.geometry.dispose();o.material.dispose();}
        if(state.arrived.length)arrivedGroup.add(pointsObject(state.arrived,0x009c68,.14,1));
        const nearestIndices=state.arrived.map(x=>{let best=Infinity,index=0;for(let i=0;i<targetCurve.length;i++){const angle=Math.acos(Math.max(-1,Math.min(1,dot(x,targetCurve[i]))));if(angle<best){best=angle;index=i;}}return{angle:best,index};});
        const meanTargetAngle=nearestIndices.length?nearestIndices.reduce((sum,item)=>sum+item.angle,0)/nearestIndices.length:0;
        const coverageBins=Array(10).fill(0);for(const item of nearestIndices)coverageBins[Math.min(9,Math.floor(10*item.index/targetCurve.length))]++;
        host.dataset.meanTargetAngle=meanTargetAngle.toFixed(4);
        host.dataset.targetCoverage=coverageBins.filter(Boolean).length;
        host.dataset.targetHistogram=coverageBins.join(',');
    }
    function clearSamples(){for(const a of state.active)a.dot.material.dispose();state.active=[];state.arrived=[];movingGroup.clear();while(pathGroup.children.length){const o=pathGroup.children.pop();o.geometry.dispose();o.material.dispose();}state.paths=[];while(arrivedGroup.children.length){const o=arrivedGroup.children.pop();o.geometry.dispose();o.material.dispose();}}
    function syncUi(){ui.train.textContent=state.running?'Pause':'Train';ui.timeVal.textContent=Number(ui.time.value).toFixed(2);fieldGroup.visible=ui.field.checked;pathGroup.visible=ui.paths.checked;sourcePoints.visible=ui.x0.checked;targetPoints.visible=ui.x1.checked;targetLine.visible=ui.x1.checked;ui.generate.disabled=model.samples.length<STRIDE*20;}
    function resize(){const w=Math.max(1,host.clientWidth),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){const dt=Math.max(0,Math.min(.05,(now-state.lastFrame)/1000));state.lastFrame=now;resize();if(ui.field.checked&&model.iterations-state.lastField>=FIELD_REFRESH)updateField();for(let i=state.active.length-1;i>=0;i--){const a=state.active[i];a.progress=Math.min(1,a.progress+dt*.42);if(a.progress>=1){movingGroup.remove(a.dot);a.dot.material.dispose();state.arrived.push(a.pts[a.pts.length-1]);state.active.splice(i,1);rebuildArrived();continue;}const f=a.progress*(a.pts.length-1),j=Math.min(a.pts.length-2,Math.floor(f)),z=f-j;const x=normalize(add(scale(a.pts[j],1-z),scale(a.pts[j+1],z)));a.dot.position.copy(world(x));}ui.stats.textContent=`steps ${model.iterations} · loss ${model.emaLoss==null?'—':model.emaLoss.toFixed(3)} · ${state.active.length} moving · ${state.arrived.length} arrived`;host.dataset.iterations=String(model.iterations);host.dataset.activeSamples=String(state.active.length);host.dataset.arrivedSamples=String(state.arrived.length);host.dataset.fieldArrows=String(fieldArrows.length);host.dataset.cameraPosition=camera.position.toArray().map(v=>v.toFixed(3)).join(',');renderer.render(scene,camera);}

    ui.train.onclick=()=>{state.running=!state.running;syncWorker();syncUi();};
    ui.reset.onclick=()=>{state.running=false;clearSamples();resetWorker();state.lastField=-FIELD_REFRESH;updateField();syncUi();};
    ui.generate.onclick=()=>addSamples(16);ui.clear.onclick=clearSamples;ui.view.onclick=()=>controls.reset();
    ui.time.oninput=()=>{updateField();syncUi();};ui.field.onchange=()=>{if(ui.field.checked)updateField();syncUi();};ui.paths.onchange=syncUi;ui.x0.onchange=syncUi;ui.x1.onchange=syncUi;
    updateField();syncUi();resize();renderer.render(scene,camera);
    const slide=window.slideIndexOf(host),renderLoop=window.makeSlideRafLoop(tick,{onStart:()=>{state.lastFrame=performance.now();resize();}});
    window.SlideAnim.register(slide,renderLoop);
    window.SlideAnim.register([slide,slide+1,slide+2],{start(){state.scopeActive=true;syncWorker();},stop(){state.scopeActive=false;syncWorker();}});
    window.SlideAnim.sync(window.slideIndexOf(document.querySelector('.slide.active')));
}