80c0b6dPlate dynamics lab: plates in realizable mounts under gravity, explicit finite differences on WebGPUJeremy Magland 1import type { ProblemSpec } from "../spec/types";
2import type { MethodDefinition } from "../methods/types";
3import type { GpuStages } from "../runner/types";
4import { generateParamsBlock } from "./params";
5import { LIBRARY_FILES, type SourceFile } from "./library";
7/** A complete, self-contained set of `.m` files plus the one to run. */
8export interface SourceBundle {
9 files: SourceFile[];
10 /** File executed as the script */
11 entry: string;
12 /** Workspace variable read back after the run */
13 resultVariable: string;
14}
16export const MAIN_FILE = "main.m";
18/**
19 * The script a run executes: the params block followed by the solver call. With
20 * `gpu` the call is written out as its three stages, since that is what the
21 * host then runs (see buildGpuStages).
22 */
23export function generateMain(
24 spec: ProblemSpec,
25 method: MethodDefinition,
26 opts: { gpu?: boolean } = {},
27): string {
28 const g = opts.gpu ? method.gpu : undefined;
29 const solve = g
30 ? [
31 `% Solve. With the GPU option the host runs ${g.setupFunction} and ${g.outputFunction} here in numbl and`,
32 `% ${g.loopFile} as its WebGPU port; without WebGPU, or when this script or ${g.loopFile} is edited,`,
33 `% all three run in numbl. The library files are in the other tabs.`,
34 `S = ${g.setupFunction}(params);`,
35 `Rec = ${g.loopFile.replace(/\.m$/, "")}(S); % replaced by the WebGPU port when the loop runs on the GPU`,
36 `result = ${g.outputFunction}(S, Rec);`,
37 ]
38 : [
39 `% Solve. ${method.solverFunction} and the functions it calls are the library files in the other tabs.`,
40 `result = ${method.solverFunction}(params);`,
41 ];
42 return [generateParamsBlock(spec), "", ...solve, ""].join("\n");
43}
45/** Bundle for running the method on a spec. */
46export function buildRunBundle(
47 spec: ProblemSpec,
48 method: MethodDefinition,
49 opts: { gpu?: boolean } = {},
50): SourceBundle {
51 return {
52 files: [{ name: MAIN_FILE, source: generateMain(spec, method, opts) }, ...LIBRARY_FILES],
53 entry: MAIN_FILE,
54 resultVariable: "result",
55 };
56}
58/**
59 * The two numbl scripts of a GPU run (see runner/executeGpu.ts): the setup,
60 * whose struct S the host reads (sparse operators as find() triplets, since
61 * a sparse matrix does not cross to the host), and the output, which takes the
62 * loop's record as the workspace variable Rec. Null for a method without a port.
63 */
64export function buildGpuStages(spec: ProblemSpec, method: MethodDefinition): GpuStages | null {
65 const g = method.gpu;
66 if (!g) return null;
67 const params = generateParamsBlock(spec);
68 const setup = [
69 params,
70 "",
71 `% Stage 1 of a GPU run: the setup, read by the WebGPU port of ${g.loopFile}.`,
72 `S = ${g.setupFunction}(params);`,
73 "% the sparse operators cross to the host as triplets",
74 "[ki, kj, kv] = find(S.KB); S.KB_i = ki; S.KB_j = kj; S.KB_v = kv;",
75 "[ki, kj, kv] = find(S.KG); S.KG_i = ki; S.KG_j = kj; S.KG_v = kv;",
76 "[ki, kj, kv] = find(S.Gs); S.Gs_i = ki; S.Gs_j = kj; S.Gs_v = kv;",
77 "[ki, kj, kv] = find(S.WB); S.WB_i = ki; S.WB_j = kj; S.WB_v = kv;",
78 "[ki, kj, kv] = find(S.WG); S.WG_i = ki; S.WG_j = kj; S.WG_v = kv;",
79 "[ki, kj, kv] = find(S.WL); S.WL_i = ki; S.WL_j = kj; S.WL_v = kv;",
80 "S = rmfield(S, 'mdl'); S = rmfield(S, 'KB'); S = rmfield(S, 'KG'); S = rmfield(S, 'Lm');",
81 "S = rmfield(S, 'Gs'); S = rmfield(S, 'WB'); S = rmfield(S, 'WG'); S = rmfield(S, 'WL');",
82 "",
83 ].join("\n");
84 const output = [
85 params,
86 "",
87 `% Stage 3 of a GPU run: Rec is the record of the WebGPU loop, passed in by the host.`,
88 `S = ${g.setupFunction}(params);`,
89 `result = ${g.outputFunction}(S, Rec);`,
90 "",
91 ].join("\n");
92 return { setup, output, loopFile: g.loopFile };
93}
95/** Bundle that evaluates the cost plan. */
96export function buildPlanBundle(spec: ProblemSpec, method: MethodDefinition): SourceBundle | null {
97 if (!method.planFunction) return null;
98 const main = [generateParamsBlock(spec), "", `plan = ${method.planFunction}(params);`, ""].join(
99 "\n",
100 );
101 return {
102 files: [{ name: "plan.m", source: main }, ...LIBRARY_FILES],
103 entry: "plan.m",
104 resultVariable: "plan",
105 };
106}
108/**
109 * Replace files in a bundle with edited versions, keeping everything else. Used
110 * when the user edits source in the code view: the run then uses exactly the
111 * text on screen.
112 */
113export function withEdits(bundle: SourceBundle, edits: Record<string, string>): SourceBundle {
114 return {
115 ...bundle,
116 files: bundle.files.map((f) =>
117 edits[f.name] !== undefined ? { ...f, source: edits[f.name] } : f,
118 ),
119 };
120}