/ concept-collection / turing-surface-cache
Sign in
concept-collection / turing-surface-cache
turing-surface-cache / src / cache / h5file.ts
231 lines · 8.8 KBBlameHistoryRaw
1/**
2 * Cache files are HDF5, in the layout of turing-surface's reference files
3 * (docs/ellipsoid-reference-spec.md there) extended with the cache's own
4 * identity at the root: the canonical spec JSON that was hashed into the
5 * object name, the app name, and the format version. A cache file is thereby
6 * also a valid reference file — turing-surface's "Compare to reference…" mode
7 * opens one as-is — and readable from Python with h5py.
8 *
9 * / attrs: app, format_version, spec_json, model, species,
10 * created_utc, adapter
11 * /backend attrs: adapter, runtime, precision
12 * /spec attrs: geometry, lmax, seed, steps, niter, lam3, t_end
13 * /spec/params attrs: a, b, D1, D2, dt
14 * /spec/geometry_params attrs: the geometry's params
15 * /grid attrs: lmax, mmax, nlat, nphi, nlm
16 * /geometry Gx, Gy, Gz float32[2*nlm]
17 * /initial one dataset per species (U, V) float32[2*nlm]
18 * /final one dataset per species (U, V) float32[2*nlm]
19 *
20 * h5wasm's browser build carries the whole HDF5 library as embedded wasm
21 * (~4 MB), so it is imported dynamically and only here: the page pays for it
22 * on the first cache hit or upload, never on startup.
23 */
24import type { ShtConfig } from '../sht/layout.ts';
25import { nlmCalc } from '../sht/layout.ts';
26import { APP_NAME, FORMAT_VERSION, canonicalJson, stepsFor, type CacheSpec } from './spec.ts';
28export interface CacheFileData {
29 spec: CacheSpec;
30 grid: ShtConfig;
31 /** Spectral state names, in order — Schnakenberg's ['U', 'V']. */
32 species: string[];
33 /** The band-limited surface's own coefficients, [re, im] per (l, m). */
34 geometry: { X: Float32Array; Y: Float32Array; Z: Float32Array };
35 /** Spectral state at t = 0 (immediately after seeding). */
36 initial: Record<string, Float32Array>;
37 /** The same, at the spec's end time. */
38 final: Record<string, Float32Array>;
39 /** Provenance: which GPU computed it. */
40 adapter: string;
43interface H5Module {
44 ready: Promise<unknown>;
45 File: new (path: string, mode: string) => H5WFile;
46 FS?: unknown;
49interface H5Attr {
50 value: unknown;
53interface H5Obj {
54 attrs: Record<string, H5Attr>;
55 get(name: string): unknown;
56 create_group(name: string): unknown;
57 create_attribute(name: string, data: unknown, shape?: unknown, dtype?: unknown): void;
58 create_dataset(args: { name: string; data: unknown; shape?: number[]; dtype?: string }): void;
61interface H5WFile extends H5Obj {
62 close(): void;
65interface EmFS {
66 writeFile(path: string, data: Uint8Array): void;
67 readFile(path: string): Uint8Array;
68 unlink(path: string): void;
71let scratchCounter = 0;
73async function withH5<T>(fn: (h5: H5Module, fs: EmFS) => T | Promise<T>): Promise<T> {
74 const h5 = (await import('h5wasm')) as unknown as H5Module;
75 const { FS } = (await h5.ready) as { FS: EmFS };
76 return fn(h5, FS);
79const groupOf = (node: H5Obj, name: string): H5Obj => {
80 const g = node.get(name) as H5Obj | null;
81 if (!g || typeof g.get !== 'function') throw new Error(`no '${name}/' group`);
82 return g;
83};
85const coeffsOf = (group: H5Obj, groupName: string, name: string, nlm: number): Float32Array => {
86 const v = (group.get(name) as { value?: unknown } | null)?.value;
87 if (!(v instanceof Float32Array)) {
88 throw new Error(`'${groupName}/${name}' is not a float32 dataset`);
89 }
90 if (v.length !== 2 * nlm) {
91 throw new Error(`'${groupName}/${name}' has ${v.length} values, expected 2*nlm = ${2 * nlm}`);
92 }
93 return v;
94};
96/** Serialize one solution to HDF5 bytes. */
97export function encodeCacheFile(data: CacheFileData): Promise<Uint8Array> {
98 return withH5((h5, FS) => {
99 const path = `/encode-${scratchCounter++}.h5`;
100 const file = new h5.File(path, 'w');
101 try {
102 const { spec, grid } = data;
103 file.create_attribute('app', APP_NAME);
104 file.create_attribute('format_version', FORMAT_VERSION);
105 file.create_attribute('spec_json', canonicalJson(spec));
106 file.create_attribute('model', spec.model);
107 file.create_attribute('species', data.species);
108 file.create_attribute('created_utc', new Date().toISOString());
109 file.create_attribute('adapter', data.adapter);
111 file.create_group('backend');
112 const backend = groupOf(file, 'backend');
113 backend.create_attribute('adapter', data.adapter);
114 backend.create_attribute('runtime', 'browser-webgpu');
115 backend.create_attribute('precision', 'fp32');
117 file.create_group('spec');
118 const specGroup = groupOf(file, 'spec');
119 specGroup.create_attribute('geometry', spec.geometry);
120 specGroup.create_attribute('lmax', spec.lmax);
121 specGroup.create_attribute('seed', spec.seed);
122 specGroup.create_attribute('steps', stepsFor(spec));
123 specGroup.create_attribute('niter', spec.niter);
124 specGroup.create_attribute('lam3', spec.lam3);
125 specGroup.create_attribute('t_end', spec.tEnd);
126 specGroup.create_group('params');
127 const params = groupOf(specGroup, 'params');
128 for (const [k, v] of Object.entries(spec.params)) params.create_attribute(k, v);
129 specGroup.create_group('geometry_params');
130 const gparams = groupOf(specGroup, 'geometry_params');
131 for (const [k, v] of Object.entries(spec.geometryParams)) gparams.create_attribute(k, v);
133 file.create_group('grid');
134 const gridGroup = groupOf(file, 'grid');
135 const nlm = nlmCalc(grid.lmax, grid.mmax);
136 gridGroup.create_attribute('lmax', grid.lmax);
137 gridGroup.create_attribute('mmax', grid.mmax);
138 gridGroup.create_attribute('nlat', grid.nlat);
139 gridGroup.create_attribute('nphi', grid.nphi);
140 gridGroup.create_attribute('nlm', nlm);
142 file.create_group('geometry');
143 const geom = groupOf(file, 'geometry');
144 geom.create_dataset({ name: 'Gx', data: data.geometry.X });
145 geom.create_dataset({ name: 'Gy', data: data.geometry.Y });
146 geom.create_dataset({ name: 'Gz', data: data.geometry.Z });
148 for (const [groupName, states] of [
149 ['initial', data.initial],
150 ['final', data.final],
151 ] as const) {
152 file.create_group(groupName);
153 const g = groupOf(file, groupName);
154 for (const name of data.species) {
155 const coeffs = states[name];
156 if (!coeffs) throw new Error(`missing ${groupName} state '${name}'`);
157 if (coeffs.length !== 2 * nlm) {
158 throw new Error(`${groupName}/${name}: ${coeffs.length} values, expected ${2 * nlm}`);
159 }
160 g.create_dataset({ name, data: coeffs });
161 }
162 }
163 } finally {
164 file.close();
165 }
166 const bytes = FS.readFile(path);
167 FS.unlink(path);
168 return bytes;
169 });
172export interface DecodedCacheFile {
173 /** Parsed from the file's own spec_json — the identity it was stored under. */
174 spec: CacheSpec;
175 species: string[];
176 initial: Record<string, Float32Array>;
177 final: Record<string, Float32Array>;
178 /** Provenance, when recorded. */
179 adapter: string;
180 created: string;
183/**
184 * Read cache-file bytes back. `expectSpecJson` is the canonical JSON this
185 * client asked the cache for; a mismatch with the file's own means the object
186 * store handed back something other than what the key promised (corruption,
187 * or a stale format), and is an error rather than a silent wrong answer.
188 */
189export function decodeCacheFile(
190 bytes: Uint8Array,
191 expectSpecJson: string,
192 expectSpecies: string[],
193): Promise<DecodedCacheFile> {
194 return withH5((h5, FS) => {
195 const path = `/decode-${scratchCounter++}.h5`;
196 FS.writeFile(path, bytes);
197 const file = new h5.File(path, 'r');
198 try {
199 const attr = (name: string): unknown => file.attrs[name]?.value;
200 const app = String(attr('app') ?? '');
201 if (app !== APP_NAME) throw new Error(`not a ${APP_NAME} file (app='${app}')`);
202 const version = Number(attr('format_version'));
203 if (version !== FORMAT_VERSION) throw new Error(`format version ${version}, expected ${FORMAT_VERSION}`);
204 const specJson = String(attr('spec_json') ?? '');
205 if (specJson !== expectSpecJson) {
206 throw new Error('file spec does not match the requested spec');
207 }
208 const spec = JSON.parse(specJson) as CacheSpec;
209 const nlm = nlmCalc(spec.lmax, spec.lmax);
210 const initialGroup = groupOf(file, 'initial');
211 const finalGroup = groupOf(file, 'final');
212 const initial: Record<string, Float32Array> = {};
213 const final: Record<string, Float32Array> = {};
214 for (const name of expectSpecies) {
215 initial[name] = coeffsOf(initialGroup, 'initial', name, nlm);
216 final[name] = coeffsOf(finalGroup, 'final', name, nlm);
217 }
218 return {
219 spec,
220 species: expectSpecies,
221 initial,
222 final,
223 adapter: String(attr('adapter') ?? ''),
224 created: String(attr('created_utc') ?? ''),
225 };
226 } finally {
227 file.close();
228 FS.unlink(path);
229 }
230 });
moveopenescclose