2dedc35turing-sphere: reaction-diffusion on the sphere, spectral solver on WebGPUJeremy Magland 1/**
61e12f1Write the solver in MATLAB and compile it to WebGPUJeremy Magland 2 * Reference solver — NOT what the app runs.
3 *
4 * The app executes the .m models under `models/` on the GPU (see `src/mgpu/`).
5 * This TypeScript port remains as an independent implementation of the same
6 * scheme, which is what makes it usable as the test oracle: `test/mgpuChecks.ts`
7 * runs both from the same seed and compares. Keep the two in step.
8 *
2dedc35turing-sphere: reaction-diffusion on the sphere, spectral solver on WebGPUJeremy Magland 9 * IMEX Euler reaction-diffusion timestepper on the sphere, ported from
10 * websph's SphericalReactionDiffusion.m. Diffusion is implicit and diagonal
11 * in spherical-harmonic space (Laplace-Beltrami eigenvalues -l(l+1));
12 * reaction is explicit on the grid:
13 *
14 * (I - dt*D_k*lap_s) u_k^{n+1} = u_k^n + dt*f_k(u^n)
15 */
16import type { ShtBackend } from './backend.ts';
17import type { ModelSpec, Params } from './models.ts';
18import { lmIndex } from '../sht/layout.ts';
20/** Grid sizes for a given lmax, dealiased for a reaction of degree pdeg
21 * (see websph README): nlat >= ((pdeg+1)*lmax+1)/2, nlon >= (pdeg+1)*lmax+1.
22 * nphi is rounded up to a power of two to keep the GPU FFT path. */
23export function gridForLmax(lmax: number, pdeg: number): { nlat: number; nphi: number } {
24 const minLat = Math.max(lmax + 1, ((pdeg + 1) * lmax + 1) / 2);
25 const nlat = 2 * Math.ceil(minLat / 2);
26 let nphi = 1;
27 while (nphi < (pdeg + 1) * lmax + 1) nphi *= 2;
28 return { nlat, nphi };
29}
31/** Seeded normal deviates: mulberry32 + Box-Muller. */
32export function makeRandn(seed: number): () => number {
33 let s = seed >>> 0;
34 const rand = () => {
35 s = (s + 0x6d2b79f5) >>> 0;
36 let t = s;
37 t = Math.imul(t ^ (t >>> 15), t | 1);
38 t ^= t + Math.imul(t ^ (t >>> 7), t | 61);
39 return ((t ^ (t >>> 14)) >>> 0) / 4294967296;
40 };
41 let spare: number | null = null;
42 return () => {
43 if (spare !== null) {
44 const v = spare;
45 spare = null;
46 return v;
47 }
48 let u = 0;
49 while (u === 0) u = rand();
50 const r = Math.sqrt(-2 * Math.log(u));
51 const th = 2 * Math.PI * rand();
52 spare = r * Math.sin(th);
53 return r * Math.cos(th);
54 };
55}
57export class Simulation {
58 readonly backend: ShtBackend;
59 readonly model: ModelSpec;
60 readonly params: Params;
61 readonly nspecies: number;
63 /** Spectral state, one interleaved-complex Float64Array (2*nlm) per species. */
64 U: Float64Array[];
65 /** Grid values per species as of the START of the last step (one step
66 * behind U; recomputed as the first stage of the next step). */
67 V: (Float32Array | Float64Array)[];
68 t = 0;
69 stepCount = 0;
71 /** Cartesian coordinates of the grid points (nlat*nphi each). */
72 readonly x: Float64Array;
73 readonly y: Float64Array;
74 readonly z: Float64Array;
75 /** Laplace-Beltrami eigenvalues l*(l+1) per spectral index (length nlm). */
76 readonly lam: Float64Array;
78 #R: Float64Array[]; // reaction scratch, one grid array per species
79 #m0Imag: number[]; // interleaved-array positions of m=0 imaginary parts
81 constructor(backend: ShtBackend, model: ModelSpec, params: Params) {
82 this.backend = backend;
83 this.model = model;
84 this.params = params;
85 this.nspecies = model.species.length;
87 const { lmax, mmax, nlat, nphi } = backend.cfg;
88 const npts = nlat * nphi;
89 this.x = new Float64Array(npts);
90 this.y = new Float64Array(npts);
91 this.z = new Float64Array(npts);
92 for (let i = 0; i < nlat; i++) {
93 const ct = backend.cosTheta[i];
94 const st = Math.sqrt(Math.max(0, 1 - ct * ct));
95 for (let j = 0; j < nphi; j++) {
96 const phi = (2 * Math.PI * j) / nphi;
97 const idx = i * nphi + j;
98 this.x[idx] = st * Math.cos(phi);
99 this.y[idx] = st * Math.sin(phi);
100 this.z[idx] = ct;
101 }
102 }
104 this.lam = new Float64Array(backend.nlm);
105 for (let m = 0; m <= mmax; m++) {
106 for (let l = m; l <= lmax; l++) {
107 this.lam[lmIndex(lmax, l, m)] = l * (l + 1);
108 }
109 }
110 this.#m0Imag = [];
111 for (let l = 0; l <= lmax; l++) {
112 this.#m0Imag.push(2 * lmIndex(lmax, l, 0) + 1);
113 }
115 this.U = [];
116 this.V = [];
117 this.#R = [];
118 for (let k = 0; k < this.nspecies; k++) {
119 this.U.push(new Float64Array(2 * backend.nlm));
120 this.V.push(new Float64Array(npts));
121 this.#R.push(new Float64Array(npts));
122 }
123 }
125 /** Project the initial conditions (band-limiting the seed noise). */
126 async init(seed: number): Promise<void> {
127 const grids = this.#R;
128 this.model.init(this.params, this.x, this.y, this.z, makeRandn(seed), grids);
129 for (let k = 0; k < this.nspecies; k++) {
130 const q = await this.backend.analys(grids[k]);
131 this.U[k].set(q);
132 this.#cleanM0(this.U[k]);
133 this.V[k] = await this.backend.synth(this.U[k]);
134 }
135 this.t = 0;
136 this.stepCount = 0;
137 }
139 /** One IMEX Euler step. */
140 async step(): Promise<void> {
141 const dt = this.params.dt;
142 const D = this.model.diffusivities(this.params);
144 // Evaluate every species on the grid before reacting any of them
145 for (let k = 0; k < this.nspecies; k++) {
146 this.V[k] = await this.backend.synth(this.U[k]);
147 }
149 this.model.reaction(this.params, this.t, this.x, this.y, this.z, this.V, this.#R);
151 for (let k = 0; k < this.nspecies; k++) {
152 const Rlm = await this.backend.analys(this.#R[k]);
153 const U = this.U[k];
154 const dD = dt * D[k];
155 for (let i = 0; i < this.backend.nlm; i++) {
156 const fac = 1 / (1 + dD * this.lam[i]);
157 U[2 * i] = (U[2 * i] + dt * Rlm[2 * i]) * fac;
158 U[2 * i + 1] = (U[2 * i + 1] + dt * Rlm[2 * i + 1]) * fac;
159 }
160 this.#cleanM0(U);
161 }
163 this.t += dt;
164 this.stepCount++;
165 }
167 /** m=0 coefficients of a real field are purely real; drop numerical junk. */
168 #cleanM0(U: Float64Array): void {
169 for (const p of this.#m0Imag) U[p] = 0;
170 }
171}