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