1/**
2 * One initial condition, on every variant's grid.
3 *
4 * The host's seeded perturbation is one normal deviate per *grid point*
5 * (src/mgpu/noise.ts), so two sessions at different lmax seeded from the same
6 * integer do not start from the same field — they start from unrelated fields
7 * that merely share a random seed. Comparing them would compare two different
8 * problems, and every number the comparison produced would be meaningless.
9 *
10 * So the field is built once, band-limited at the *coarsest* variant's lmax,
11 * and evaluated on each variant's own grid:
12 *
13 * 1. white noise on the coarsest grid
14 * 2. analysed there -> coefficients up to lmax_min
15 * 3. zero-padded into each variant's coefficient layout
16 * 4. synthesized on that variant's grid
17 *
18 * Steps 3 and 4 are exact: the field is band-limited at lmax_min, and every
19 * variant's band contains that, so each one receives the *same function*
20 * sampled where it needs it. Running each model's own `init` on it then leaves
21 * every session holding the identical spectral state (zero-padded), which is
22 * what makes a pointwise comparison at later times mean something.
23 *
24 * The coarsest variant gets the projected field too, not the raw white noise
25 * it was analysed from — otherwise it alone would start somewhere slightly
26 * different from the others.
27 */
28import { lmIndex, nlmCalc } from '../sht/layout.ts';
29import { seededNoise } from '../mgpu/noise.ts';
30import type { ModelSession } from '../mgpu/session.ts';
32/**
33 * Re-index coefficients from a band limit into a wider one's layout, zero-
34 * filling the degrees the source does not have. Both layouts are SHTNS
35 * m-major with mmax = lmax, so nothing but the index mapping changes.
36 */
37export function prolongCoeffs(
38 q: Float32Array,
39 lmaxFrom: number,
40 lmaxTo: number,
41): Float32Array {
42 if (lmaxTo === lmaxFrom) return q;
43 if (lmaxTo < lmaxFrom) {
44 throw new Error(`prolongCoeffs: cannot widen ${lmaxFrom} into a smaller ${lmaxTo}`);
45 }
46 const out = new Float32Array(2 * nlmCalc(lmaxTo, lmaxTo));
47 for (let m = 0; m <= lmaxFrom; m++) {
48 for (let l = m; l <= lmaxFrom; l++) {
49 const from = 2 * lmIndex(lmaxFrom, l, m);
50 const to = 2 * lmIndex(lmaxTo, l, m);
51 out[to] = q[from];
52 out[to + 1] = q[from + 1];
53 }
54 }
55 return out;
56}
58/**
59 * The same band-limited perturbation, sampled on each session's grid. Order
60 * follows `sessions`. Nothing may be in flight on any session's transform
61 * plan — the one-off analys/synth here use the plan's own scratch buffers.
62 */
63export async function sharedNoise(
64 sessions: ModelSession[],
65 amp: number,
66 seed: number,
67): Promise<Float32Array[]> {
68 let base = sessions[0];
69 for (const s of sessions) if (s.cfg.lmax < base.cfg.lmax) base = s;
70 const coeffs = await base.sht.analys(seededNoise(base.npts, amp, seed));
71 const out: Float32Array[] = [];
72 for (const s of sessions) {
73 out.push(await s.sht.synth(prolongCoeffs(coeffs, base.cfg.lmax, s.cfg.lmax)));
74 }
75 return out;
76}