WebGPU PINN Training Studio

// tiny PINN forward with 2nd-order forward-mode autodiff
export type Act = 'tanh'|'sin'|'swish';
function act(z: number, kind: Act){
  if(kind==='tanh'){ const f=Math.tanh(z); return {f, fp:1-f*f, fpp:-2*f*(1-f*f)}; }
  if(kind==='sin'){ return {f:Math.sin(z), fp:Math.cos(z), fpp:-Math.sin(z)}; }
  const s=1/(1+Math.exp(-z)); const f=z*s; const ds=s*(1-s);
  const fp=s+z*ds; const d2=ds*(1-2*s); const fpp=2*ds+z*d2; return {f, fp, fpp};
}
export function forward(x:number,t:number, layers:{W:number[][],b:number[]}[], Wout:number[], bout:number, kind:Act){
  // a0=[x,t], da/dx=[1,0] etc.
  let a=[x,t]; let dadx=[1,0], dadt=[0,1], d2adx2=[0,0], d2adt2=[0,0];
  const cache=[];
  for(const {W,b} of layers){
    const out=W[0].length, z=new Array(out).fill(0), dzdx=new Array(out).fill(0), dzdt=new Array(out).fill(0), d2zdx2=new Array(out).fill(0), d2zdt2=new Array(out).fill(0);
    for(let j=0;j<out;j++){ let s=b[j]; let sx=0, st=0, sxx=0, stt=0;
      for(let k=0;k<a.length;k++){ s+=a[k]*W[k][j]; sx+=dadx[k]*W[k][j]; st+=dadt[k]*W[k][j]; sxx+=d2adx2[k]*W[k][j]; stt+=d2adt2[k]*W[k][j]; }
      z[j]=s; dzdx[j]=sx; dzdt[j]=st; d2zdx2[j]=sxx; d2zdt2[j]=stt;
    }
    const an=new Array(out), dadxn=new Array(out), dadtn=new Array(out), d2adx2n=new Array(out), d2adt2n=new Array(out);
    for(let j=0;j<out;j++){ const {f,fp,fpp}=act(z[j],kind);
      an[j]=f; dadxn[j]=fp*dzdx[j]; dadtn[j]=fp*dzdt[j];
      d2adx2n[j]=fpp*dzdx[j]*dzdx[j]+fp*d2zdx2[j];
      d2adt2n[j]=fpp*dzdt[j]*dzdt[j]+fp*d2zdt2[j];
    }
    cache.push({z, a}); a=an; dadx=dadxn; dadt=dadtn; d2adx2=d2adx2n; d2adt2=d2adt2n;
  }
  let u=bout, ut=0, ux=0, uxx=0, utt=0;
  for(let k=0;k<a.length;k++){ u+=a[k]*Wout[k]; ut+=dadt[k]*Wout[k]; ux+=dadx[k]*Wout[k]; uxx+=d2adx2[k]*Wout[k]; utt+=d2adt2[k]*Wout[k]; }
  return {u, ut, ux, uxx, utt, a_last:a, dadx, dadt, d2adx2, d2adt2, cache};
}
export function residual(f:ReturnType<typeof forward>, pde:'heat'|'burgers'|'wave', p:{alpha:number,nu:number,c:number}){
  if(pde==='heat') return f.ut - p.alpha*f.uxx;
  if(pde==='wave') return f.utt - p.c*p.c*f.uxx;
  return f.ut + f.u*f.ux - p.nu*f.uxx; // burgers
}
export function pinnLoss(colloc:{x:number,t:number}[], data:{x:number,t:number,target:number}[], model:any, pde:any, params:any, lambda=1){
  let Lp=0; for(const pt of colloc){ const fr=forward(pt.x,pt.t,model.layers,model.Wout,model.bout,model.act); Lp+= residual(fr,pde,params)**2; } Lp/=colloc.length;
  let Ld=0; for(const pt of data){ const fr=forward(pt.x,pt.t,model.layers,model.Wout,model.bout,model.act); Ld+=(fr.u-pt.target)**2; } Ld/=data.length;
  return {total: Ld+lambda*Lp, Lp, Ld};
}