concept-collection / turing-surface-cache
turing-surface-cache / src / cache / h5file.ts
268 lines · 10.4 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;
41 /** And what was driving it — a browser, or the command line's Dawn. */
42 runtime: string;
45interface H5Module {
46 ready: Promise<unknown>;
47 File: new (path: string, mode: string) => H5WFile;
48 FS?: unknown;
51interface H5Attr {
52 value: unknown;
55interface H5Obj {
56 attrs: Record<string, H5Attr>;
57 get(name: string): unknown;
58 create_group(name: string): unknown;
59 create_attribute(name: string, data: unknown, shape?: unknown, dtype?: unknown): void;
60 create_dataset(args: { name: string; data: unknown; shape?: number[]; dtype?: string }): void;
63interface H5WFile extends H5Obj {
64 close(): void;
67interface EmFS {
68 writeFile(path: string, data: Uint8Array): void;
69 readFile(path: string): Uint8Array;
70 unlink(path: string): void;
73let scratchCounter = 0;
75/**
76 * Where the scratch file that h5wasm reads or writes lives.
77 *
78 * In the browser it lives in h5wasm's own in-memory filesystem, where any
79 * absolute path will do and nothing touches a disk. The node build is
80 * compiled with NODERAWFS, which is to say its filesystem *is* the real one:
81 * the same path would name a file in the root directory, which fails with a
82 * wall of HDF5 diagnostics rather than an error anyone could act on. A real
83 * temporary directory is therefore used there, and the command line replaces
84 * this default with the platform's own (src/cli/fill.ts).
85 */
86let scratchDir = __NODE_BUILD__ ? '/tmp/' : '/';
88/** Set the directory for those scratch files; node only. */
89export function setScratchDir(dir: string): void {
90 scratchDir = dir.endsWith('/') ? dir : `${dir}/`;
93/** Two fills on one machine share that real directory; a page's filesystem is
94 * its own, so there is nothing to distinguish there. */
95const scratchTag = __NODE_BUILD__ ? `${process.pid}-` : '';
97const scratchPath = (what: string): string =>
98 `${scratchDir}turing-surface-cache-${scratchTag}${what}-${scratchCounter++}.h5`;
100async function withH5<T>(fn: (h5: H5Module, fs: EmFS) => T | Promise<T>): Promise<T> {
101 // Two builds of the same library: the browser one carries the wasm inside
102 // the bundle, the node one reads it off disk. __NODE_BUILD__ is a build-time
103 // constant (see vite.config.ts and vite.cli.config.ts), so whichever branch
104 // this build is not takes no part in it.
105 const h5 = (await (__NODE_BUILD__
106 ? import('h5wasm/node')
107 : import('h5wasm'))) as unknown as H5Module;
108 const { FS } = (await h5.ready) as { FS: EmFS };
109 return fn(h5, FS);
112const groupOf = (node: H5Obj, name: string): H5Obj => {
113 const g = node.get(name) as H5Obj | null;
114 if (!g || typeof g.get !== 'function') throw new Error(`no '${name}/' group`);
115 return g;
116};
118const coeffsOf = (group: H5Obj, groupName: string, name: string, nlm: number): Float32Array => {
119 const v = (group.get(name) as { value?: unknown } | null)?.value;
120 if (!(v instanceof Float32Array)) {
121 throw new Error(`'${groupName}/${name}' is not a float32 dataset`);
122 }
123 if (v.length !== 2 * nlm) {
124 throw new Error(`'${groupName}/${name}' has ${v.length} values, expected 2*nlm = ${2 * nlm}`);
125 }
126 return v;
127};
129/** Serialize one solution to HDF5 bytes. */
130export function encodeCacheFile(data: CacheFileData): Promise<Uint8Array> {
131 return withH5((h5, FS) => {
132 const path = scratchPath('encode');
133 const file = new h5.File(path, 'w');
134 try {
135 const { spec, grid } = data;
136 file.create_attribute('app', APP_NAME);
137 file.create_attribute('format_version', FORMAT_VERSION);
138 file.create_attribute('spec_json', canonicalJson(spec));
139 file.create_attribute('model', spec.model);
140 file.create_attribute('species', data.species);
141 file.create_attribute('created_utc', new Date().toISOString());
142 file.create_attribute('adapter', data.adapter);
144 file.create_group('backend');
145 const backend = groupOf(file, 'backend');
146 backend.create_attribute('adapter', data.adapter);
147 backend.create_attribute('runtime', data.runtime);
148 backend.create_attribute('precision', 'fp32');
150 file.create_group('spec');
151 const specGroup = groupOf(file, 'spec');
152 specGroup.create_attribute('geometry', spec.geometry);
153 specGroup.create_attribute('lmax', spec.lmax);
154 specGroup.create_attribute('seed', spec.seed);
155 specGroup.create_attribute('steps', stepsFor(spec));
156 specGroup.create_attribute('niter', spec.niter);
157 specGroup.create_attribute('lam3', spec.lam3);
158 specGroup.create_attribute('t_end', spec.tEnd);
159 specGroup.create_group('params');
160 const params = groupOf(specGroup, 'params');
161 for (const [k, v] of Object.entries(spec.params)) params.create_attribute(k, v);
162 specGroup.create_group('geometry_params');
163 const gparams = groupOf(specGroup, 'geometry_params');
164 for (const [k, v] of Object.entries(spec.geometryParams)) gparams.create_attribute(k, v);
166 file.create_group('grid');
167 const gridGroup = groupOf(file, 'grid');
168 const nlm = nlmCalc(grid.lmax, grid.mmax);
169 gridGroup.create_attribute('lmax', grid.lmax);
170 gridGroup.create_attribute('mmax', grid.mmax);
171 gridGroup.create_attribute('nlat', grid.nlat);
172 gridGroup.create_attribute('nphi', grid.nphi);
173 gridGroup.create_attribute('nlm', nlm);
175 file.create_group('geometry');
176 const geom = groupOf(file, 'geometry');
177 geom.create_dataset({ name: 'Gx', data: data.geometry.X });
178 geom.create_dataset({ name: 'Gy', data: data.geometry.Y });
179 geom.create_dataset({ name: 'Gz', data: data.geometry.Z });
181 for (const [groupName, states] of [
182 ['initial', data.initial],
183 ['final', data.final],
184 ] as const) {
185 file.create_group(groupName);
186 const g = groupOf(file, groupName);
187 for (const name of data.species) {
188 const coeffs = states[name];
189 if (!coeffs) throw new Error(`missing ${groupName} state '${name}'`);
190 if (coeffs.length !== 2 * nlm) {
191 throw new Error(`${groupName}/${name}: ${coeffs.length} values, expected ${2 * nlm}`);
192 }
193 g.create_dataset({ name, data: coeffs });
194 }
195 }
196 } finally {
197 file.close();
198 }
199 try {
200 return FS.readFile(path);
201 } finally {
202 // Under NODERAWFS this is a real file in a real temporary directory, so
203 // it is removed on the way out however this ends.
204 FS.unlink(path);
205 }
206 });
209export interface DecodedCacheFile {
210 /** Parsed from the file's own spec_json — the identity it was stored under. */
211 spec: CacheSpec;
212 species: string[];
213 initial: Record<string, Float32Array>;
214 final: Record<string, Float32Array>;
215 /** Provenance, when recorded. */
216 adapter: string;
217 created: string;
220/**
221 * Read cache-file bytes back. `expectSpecJson` is the canonical JSON this
222 * client asked the cache for; a mismatch with the file's own means the object
223 * store handed back something other than what the key promised (corruption,
224 * or a stale format), and is an error rather than a silent wrong answer.
225 */
226export function decodeCacheFile(
227 bytes: Uint8Array,
228 expectSpecJson: string,
229 expectSpecies: string[],
230): Promise<DecodedCacheFile> {
231 return withH5((h5, FS) => {
232 const path = scratchPath('decode');
233 FS.writeFile(path, bytes);
234 const file = new h5.File(path, 'r');
235 try {
236 const attr = (name: string): unknown => file.attrs[name]?.value;
237 const app = String(attr('app') ?? '');
238 if (app !== APP_NAME) throw new Error(`not a ${APP_NAME} file (app='${app}')`);
239 const version = Number(attr('format_version'));
240 if (version !== FORMAT_VERSION) throw new Error(`format version ${version}, expected ${FORMAT_VERSION}`);
241 const specJson = String(attr('spec_json') ?? '');
242 if (specJson !== expectSpecJson) {
243 throw new Error('file spec does not match the requested spec');
244 }
245 const spec = JSON.parse(specJson) as CacheSpec;
246 const nlm = nlmCalc(spec.lmax, spec.lmax);
247 const initialGroup = groupOf(file, 'initial');
248 const finalGroup = groupOf(file, 'final');
249 const initial: Record<string, Float32Array> = {};
250 const final: Record<string, Float32Array> = {};
251 for (const name of expectSpecies) {
252 initial[name] = coeffsOf(initialGroup, 'initial', name, nlm);
253 final[name] = coeffsOf(finalGroup, 'final', name, nlm);
254 }
255 return {
256 spec,
257 species: expectSpecies,
258 initial,
259 final,
260 adapter: String(attr('adapter') ?? ''),
261 created: String(attr('created_utc') ?? ''),
262 };
263 } finally {
264 file.close();
265 FS.unlink(path);
266 }
267 });