From 431499b8e2683d8c62ea1ef6a8a8144799bb02e5 Mon Sep 17 00:00:00 2001 From: David Ackerman Date: Fri, 16 Oct 2026 15:23:16 -0400 Subject: [PATCH 1/2] =?UTF-8?q?Engine.materialize=20writes=20exactly=20the?= =?UTF-8?q?=20box=20when=20asked,=20as=20cellmap-flow's=20Run=20tab=20want?= =?UTF-8?q?s=20a=20region=20saved=20as=20that=20region=20and=20not=20as=20?= =?UTF-8?q?the=20whole=20level=20with=20the=20region's=20chunks=20in=20it:?= =?UTF-8?q?=20extent=20"box"=20(the=20default,=20"level",=20is=20the=20arr?= =?UTF-8?q?ay=20of=20the=20level's=20whole=20shape=20with=20the=20box's=20?= =?UTF-8?q?chunks=20rounded=20out=20to=20whole=20chunks,=20as=20before)=20?= =?UTF-8?q?makes=20the=20array=20hi=20=E2=88=92=20lo=20voxels=20(the=20box?= =?UTF-8?q?=20clipped=20to=20the=20level=20and=20rounded=20out=20to=20whol?= =?UTF-8?q?e=20voxels,=20not=20to=20chunks),=20its=20translation=20the=20l?= =?UTF-8?q?evel's=20moved=20by=20lo=20voxels=20(layout.ts:=20a=20dataset's?= =?UTF-8?q?=20lo,=20so=20voxel=200's=20centre=20lands=20on=20the=20level's?= =?UTF-8?q?=20voxel=20lo),=20its=20chunks=20the=20served=20chunk=20shape?= =?UTF-8?q?=20counted=20from=20lo,=20those=20at=20its=20far=20edge=20store?= =?UTF-8?q?d=20whole=20with=20the=20fill=20value=20past=20the=20box,=20as?= =?UTF-8?q?=20zarr=20v3's=20regular=20grid=20says=20and=20zarr=20v2's=20on?= =?UTF-8?q?e=20chunk=20shape=20implies=20(zarrita,=20zarr-python=203.4=20a?= =?UTF-8?q?nd=20tensorstore=20read=20such=20a=20store=20back,=20and=20the?= =?UTF-8?q?=20last=20two=20refuse=20an=20edge=20chunk=20cut=20short);=20wh?= =?UTF-8?q?ere=20lo=20is=20off=20the=20served=20chunks'=20edges=20each=20o?= =?UTF-8?q?f=20the=20box's=20chunks=20is=20stitched=20from=20the=20served?= =?UTF-8?q?=20chunks=20it=20overlaps,=20up=20to=20two=20along=20each=20axi?= =?UTF-8?q?s,=20and=20since=20a=20chunk=20only=20materialize=20asked=20for?= =?UTF-8?q?=20is=20not=20kept=20(ac41da8)=20each=20served=20chunk=20is=20a?= =?UTF-8?q?sked=20for=20once:=20writeChunks=20walks=20them=20in=20C=20orde?= =?UTF-8?q?r=20through=20the=20same=20background=20queue=20and=20stitcher(?= =?UTF-8?q?)=20copies=20each=20into=20the=20box's=20chunks=20it=20overlaps?= =?UTF-8?q?,=20a=20box=20chunk=20written=20once=20its=20last=20part=20is?= =?UTF-8?q?=20in=20and=20a=20buffer=20made=20only=20when=20a=20part=20hold?= =?UTF-8?q?s=20a=20value=20not=20zero,=20so=20at=20most=20a=20slab=20of=20?= =?UTF-8?q?the=20box's=20chunks=20(one=20chunk=20thick=20along=20the=20fir?= =?UTF-8?q?st=20axis=20on=20which=20lo=20is=20off=20the=20edges),=20a=20ro?= =?UTF-8?q?w=20and=20one=20more=20are=20held,=20none=20where=20lo=20is=20o?= =?UTF-8?q?n=20chunk=20edges,=20the=20served=20chunks=20then=20written=20a?= =?UTF-8?q?s=20they=20are;=20a=20box=20of=20the=20whole=20level=20is=20the?= =?UTF-8?q?=20level's=20extent=20byte=20for=20byte;=20progress=20gains=20c?= =?UTF-8?q?omputed=20and=20needed=20(the=20served=20chunks=20computed,=20o?= =?UTF-8?q?f=20those=20the=20box=20covers)=20beside=20done=20and=20total?= =?UTF-8?q?=20(the=20array's=20chunks),=20each=20only=20growing;=20a=20ser?= =?UTF-8?q?ved=20chunk=20that=20fails=20in=20a=20box=20is=20named=20as=20t?= =?UTF-8?q?he=20engine=20names=20it=20(view/level/z.y.x);=20the=20tests=20?= =?UTF-8?q?cover=20the=20box's=20metadata=20and=20translation,=20aligned,?= =?UTF-8?q?=20unaligned,=20cropped=20and=20smaller-than-a-chunk=20boxes,?= =?UTF-8?q?=20edge=20chunks,=20both=20formats=20read=20back=20by=20zarrita?= =?UTF-8?q?,=20channel=20axes,=20skipEmpty,=20abort,=20failure,=20each=20s?= =?UTF-8?q?erved=20chunk=20asked=20for=20once=20and=20the=20chunks=20held;?= =?UTF-8?q?=20the=20formats,=20caching=20and=20design=20pages=20and=20the?= =?UTF-8?q?=20CHANGELOG=20say=20so?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- CHANGELOG.md | 9 +- docs/concepts/caching.md | 7 +- docs/concepts/formats.md | 40 +++-- docs/design.md | 32 ++-- web/src/engine.ts | 19 ++- web/src/layout.ts | 53 ++++--- web/src/materialize.ts | 252 ++++++++++++++++++++++------- web/test/materialize.test.ts | 296 ++++++++++++++++++++++++++++++++++- 8 files changed, 599 insertions(+), 109 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index d04728d..5ad3125 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -153,7 +153,14 @@ The first release, an alpha: the plugin API can still change before 0.1.0. - `Engine.materialize(view, store, options)` writes a box of a served view (or a whole level) into a store as an OME-Zarr of its own, zarr v2 with OME 0.4 or zarr v3 with OME 0.5, gzipped, its empty chunks left out, the metadata the engine serves for the view - (`web/src/layout.ts`, shared with the routes) with the one level as its one dataset; + (`web/src/layout.ts`, shared with the routes) with the one level as its one dataset: the + level's whole shape, the box's chunks in it rounded out to whole chunks (`extent: + "level"`, the default), or exactly the box (`extent: "box"`), `hi − lo` voxels with the + level's translation moved by `lo` voxels, its chunks the served chunk shape counted from + `lo`, each stitched from the served chunks it overlaps (each computed once, about a slab + of the box's chunks held at a time), those at its far edge stored whole with the fill + value past the box; progress counts the array's chunks done and the served chunks + computed (`computed`, `needed`); `directoryStore()` writes into a folder picked with the File System Access API (or the origin's private file system), `memoryStore()` into memory ([formats](docs/concepts/formats.md#a-region-written-to-a-store-browser-engine)). Its diff --git a/docs/concepts/caching.md b/docs/concepts/caching.md index d44158c..7a95108 100644 --- a/docs/concepts/caching.md +++ b/docs/concepts/caching.md @@ -231,8 +231,11 @@ whole region, so the queue holds a few of its chunks at a time however large the A chunk only materialize asked for is not kept once written (the kept chunks are for clients that fetch a chunk again), so a large region does not push what the viewer is looking at out of the 256 MB; one a client asked for too, while it was on its way, is kept -as any other. A staged view's stage outputs are kept as usual (512 MB), so writing a large -region does cycle through those. +as any other. So materialize asks for each chunk once: a box written as exactly the box +(`extent: "box"`) that does not start on chunk edges has each of its chunks stitched from +the served chunks it overlaps as they come, rather than asking a served chunk again for +each of the box's chunks it overlaps. A staged view's stage outputs are kept as usual +(512 MB), so writing a large region does cycle through those. ## What has been computed diff --git a/docs/concepts/formats.md b/docs/concepts/formats.md index 47109ef..f28bd16 100644 --- a/docs/concepts/formats.md +++ b/docs/concepts/formats.md @@ -105,20 +105,40 @@ The browser engine also writes a view, or a box of it, into a store as files folder the visitor picks, the origin's private file system or memory. What it writes is the layout it serves for the view, one level of it, so the result opens wherever the view did and lands in the same place: the group lists the level written as its one dataset, -`0`, with that level's scale and translation; the level's array has the level's whole -shape (a box keeps its place in it), the chunk shape and dtype served, and chunks -compressed with gzip unless asked not to be; and the chunks of the box, rounded out to -whole chunks. A chunk of nothing but zeros is left out (zarr reads a missing chunk as the -fill value, 0), unless asked to be written. A channel axis the view's ops add is the -first axis and one chunk, as served. As zarr v2 with OME-Zarr 0.4 (the default, which -more tools read): +`0`, with that level's scale; the array has the chunk shape and dtype served, and chunks +compressed with gzip unless asked not to be. What the array covers is the `extent`: + +* `"level"` (the default): the level's whole shape, with the level's translation, so a box + keeps its place in it; the chunks of the box are written, rounded out to whole chunks, + each the chunk the engine serves. +* `"box"`: exactly the box, `hi − lo` voxels (clipped to the level, rounded out to whole + voxels, not to chunks), with the level's translation moved by `lo` voxels (the + translation is where voxel 0's centre is, so it lands where the view's voxel `lo` does). + Its chunks have the served chunk shape, counted from `lo`. Where `lo` is not on the + served chunks' edges, each of the box's chunks is stitched from the served chunks it + overlaps, up to two along each axis, each served chunk computed once: they are walked in + C order and each is copied into the box's chunks it overlaps, a box chunk written once + the last of them is in, so only the box's chunks begun and not yet complete are held, at + most a slab of them (one chunk thick along the first axis on which `lo` is off the + chunk edges), a row and one more; none where `lo` is on chunk edges, the box's chunks + then the served chunks as they are. The chunks at the box's far edge overhang it and are + stored whole, zero (the fill value) past the box, as zarr stores a chunk at an array's + edge (zarr v3's regular grid: border chunks have the full chunk size, the fill value + recommended outside the array; zarr v2: every chunk has the same shape). A box of the + whole level is the same store as the level's extent, byte for byte. + +A chunk of nothing but zeros is left out (zarr reads a missing chunk as the fill value, +0), unless asked to be written. A channel axis the view's ops add is the first axis and +one chunk, as served, whole whatever the extent. Progress counts the array's chunks done +(written or skipped) of its total, and the served chunks computed of those the box needs; +each only grows. As zarr v2 with OME-Zarr 0.4 (the default, which more tools read): ```text .zgroup {"zarr_format": 2} .zattrs {...extra attributes, "multiscales": [{"name": , "axes": [...], "datasets": [{"path": "0", "coordinateTransformations": [scale, translation]}], "version": "0.4"}]} -0/.zarray {"zarr_format": 2, "shape": , "chunks": , "dtype": ", "chunks": , "dtype": "// each chunk, gzipped ``` diff --git a/docs/design.md b/docs/design.md index e147a0e..dc1c71c 100644 --- a/docs/design.md +++ b/docs/design.md @@ -525,18 +525,28 @@ as text). What such a page gets, in the order it uses it: file system, `navigator.storage.getDirectory()`), its subdirectories made as needed and each file written whole through `createWritable()`, and `memoryStore()` one of a `Map`, for tests and for zipping a small result. The options: the view's `level` (the finest by - default), a voxel box `lo`, `hi` on it (the whole level by default; rounded out to whole - chunks and clipped to the level), `format` (`"zarr2"`, OME-Zarr 0.4, the default, since - more tools read it; or `"zarr3"`, OME-Zarr 0.5), `compression` (`"gzip"`, the default, or - `null`), `skipEmpty` (a chunk of nothing but zeros left out, by default), `concurrency`, - `signal` (an `AbortSignal`), `onProgress` (`done`, `total`, `written`, `bytes`) and + default), a voxel box `lo`, `hi` on it (the whole level by default; clipped to the level), + `extent` (`"level"`, the default: the array is the level's whole shape and the box's + chunks are written, rounded out to whole chunks; `"box"`: the array is exactly the box, + translated to its first voxel, its chunks the served chunk shape counted from `lo`, each + stitched from the served chunks it overlaps), `format` (`"zarr2"`, OME-Zarr 0.4, the + default, since more tools read it; or `"zarr3"`, OME-Zarr 0.5), `compression` (`"gzip"`, + the default, or `null`), `skipEmpty` (a chunk of nothing but zeros left out, by default), + `concurrency`, `signal` (an `AbortSignal`), `onProgress` (`done`, `total`, `written`, + `bytes` for the array's chunks; `computed`, `needed` for the served chunks) and `attributes` (the consumer's provenance, beside the OME metadata). It resolves with `{chunks, written, skipped, bytes, aborted}`. The layout is the one the engine serves for the view, from one module (`web/src/layout.ts`) for both, with the one level written as - the group's one dataset at `0` with its own scale and translation, so it lands where the - view does ([formats](concepts/formats.md#a-region-written-to-a-store-browser-engine)). A - single level rather than the pyramid: the engine downsamples nothing itself, and every - level would be the model run again on its own grid. The metadata is written first, so a + the group's one dataset at `0` with its own scale and translation (moved to the box's + first voxel, with the box's extent), so it lands where the view does + ([formats](concepts/formats.md#a-region-written-to-a-store-browser-engine)). A box that + does not start on the served chunks' edges is not computed again on the box's own grid: + the engine keeps no chunk only materialize asked for, so asking a served chunk again for + each of the box's chunks it overlaps would compute it up to eight times; instead the + served chunks are each asked for once, in C order, and copied into the box's chunks they + overlap, a box chunk written once complete, so about a slab of the box's chunks is held + at once. A single level rather than the pyramid: the engine downsamples nothing itself, + and every level would be the model run again on its own grid. The metadata is written first, so a store whose writing stopped is still an OME-Zarr, zero where no chunk was written. The chunks come from the same producer the routes use, through the same queue, but in the background: they wait behind the viewer's, in all its slots but one, and are never dropped @@ -547,8 +557,8 @@ as text). What such a page gets, in the order it uses it: and why, once the writes under way have ended. This is a first piece of [roadmap](roadmap.md) item 2: an explicit write of a box, asked for by the page, not yet the cache on disk that browsing fills. The parts that need no engine (the chunks of a - box, the files, gzip, the loop over any producer, the stores) run in `npm test`, the - result read back by zarrita (`web/test/materialize.test.ts`). + box, the files, gzip, the stitching, the loop over any producer, the stores) run in + `npm test`, the result read back by zarrita (`web/test/materialize.test.ts`). * `Engine.unserve(view)` stops serving a view: its infos, producer and kept chunks are forgotten and chunks of it waiting to be computed dropped, so a page that serves a model again under a new name (viewers cache by URL, so an edit is a new name) does not keep the diff --git a/web/src/engine.ts b/web/src/engine.ts index d6add92..f5ac1c6 100644 --- a/web/src/engine.ts +++ b/web/src/engine.ts @@ -501,14 +501,17 @@ export class Engine { /** Write a region of `view` into `store` as an OME-Zarr of its own (materialize.ts): one * level of the view (`options.level`, the finest by default) over a box of its voxels - * (`lo`, `hi`, rounded out to whole chunks; the whole level by default), as zarr v2 with - * OME-Zarr 0.4 or zarr v3 with OME-Zarr 0.5, gzipped, its chunks of nothing but zeros left - * out. The group and the array are written first, so the store reads as the view from the - * start, zero where no chunk is yet; then the chunks, each computed as a client's request - * for it is (kept, or shared with one on its way, and kept after) but in the background: - * behind every chunk a client waits for, in all the queue's slots but one, and never - * dropped because a client gave up on it (demand.ts). `concurrency` chunks are asked for - * at once, as many as the queue has slots by default. Resolves with what was written, or + * (`lo`, `hi`; the whole level by default), as zarr v2 with OME-Zarr 0.4 or zarr v3 with + * OME-Zarr 0.5, gzipped, its chunks of nothing but zeros left out. With `extent` "level" + * (the default) the array is the level's whole shape and the box's chunks are written, + * rounded out to whole chunks; with "box" it is exactly the box, translated to its first + * voxel, each of its chunks stitched from the view's it overlaps. The group and the array + * are written first, so the store reads as the view from the start, zero where no chunk is + * yet; then the chunks, each of the view's computed once, as a client's request for it is + * (kept, or shared with one on its way, and kept after) but in the background: behind + * every chunk a client waits for, in all the queue's slots but one, and never dropped + * because a client gave up on it (demand.ts). `concurrency` chunks are asked for at once, + * as many as the queue has slots by default. Resolves with what was written, or * on `signal`'s abort with `aborted` (chunks waiting for their turn are dropped, the ones * written stay); rejects naming the chunk that failed to compute or be written, once the * writes under way have ended. */ diff --git a/web/src/layout.ts b/web/src/layout.ts index ecf24d5..97d4af7 100644 --- a/web/src/layout.ts +++ b/web/src/layout.ts @@ -1,9 +1,9 @@ // A view's OME-Zarr layouts, as the engine serves them (engine.ts's routes) and writes them // into a store (materialize.ts), from one place so the two cannot differ: zarr v3 with // OME-Zarr 0.5, and zarr v2 with OME-Zarr 0.4. The group's metadata, a level's array -// metadata and a chunk's key, as objects ready for JSON. A channel axis the view's ops add -// (`added`) comes first and is one chunk, as the routes serve it. Nothing here imports a -// value, so Node's tests run it. +// metadata (or a box of the level's, which materialize writes) and a chunk's key, as +// objects ready for JSON. A channel axis the view's ops add (`added`) comes first and is +// one chunk, as the routes serve it. Nothing here imports a value, so Node's tests run it. import type { ViewAxis, ViewInfo } from "./types.ts"; export const TO_SECONDS: Record = { s: 1, second: 1, millisecond: 1e-3, ms: 1e-3, minute: 60, hour: 3600, day: 86400 }; @@ -12,27 +12,41 @@ const OME_UNIT: Record = { nm: "nanometer", um: "micrometer", m: export const V2_DTYPE: Record = { uint8: "|u1", int8: "|i1", uint16: " = { + uint8: 1, int8: 1, uint16: 2, int16: 2, uint32: 4, int32: 4, uint64: 8, int64: 8, float32: 4, float64: 8, +}; export const isTime = (a: ViewAxis) => a.unit in TO_SECONDS; /** A level of a view as one of a group's datasets: its path in the group, and which of the - * view's levels it is. */ -export interface Dataset { path: string; level: number } + * view's levels it is; `lo` (a box of the level, materialize.ts): the level's voxel that is + * the dataset's voxel 0, on the level's own axes (none: the level's voxel 0). */ +export interface Dataset { path: string; level: number; lo?: number[] } /** A zarr v3 codec after `bytes` (a compressor). */ export interface Codec { name: string; configuration: Record } +/** Where voxel `lo` of a level of view `v` has its centre (the level's origin when `lo` is + * left out): the origin shifted by `lo` voxels, on the view's own axes. */ +export const translation = (v: ViewInfo, level: number, lo?: number[]) => { + const { origin, voxel } = v.levels[level]; + return lo ? origin.map((o, a) => o + lo[a] * voxel[a]) : origin; +}; + /** The OME multiscales entry of view `name`, whose datasets are the levels `datasets` name: * the channel axes its ops add, then its own (space, or time by their unit); each level's - * scale its voxel size and its translation its origin, so it lands where the view does. */ + * scale its voxel size and its translation its origin (where voxel 0's centre is), so it + * lands where the view does; a dataset that starts at voxel `lo` of its level is translated + * by `lo` voxels more, so its voxel 0 lands on the level's voxel `lo`. */ export function multiscale(name: string, v: ViewInfo, datasets: Dataset[]) { const added = (v.added ?? []).map(() => 0); return { name, axes: [...added.map(() => ({ name: "c", type: "channel" })), ...v.axes.map((a) => ({ name: a.name, type: isTime(a) ? "time" : "space", unit: OME_UNIT[a.unit] ?? (a.unit || undefined) }))], - datasets: datasets.map(({ path, level }) => ({ path, coordinateTransformations: [ + datasets: datasets.map(({ path, level, lo }) => ({ path, coordinateTransformations: [ { type: "scale", scale: [...added.map(() => 1), ...v.levels[level].voxel] }, - { type: "translation", translation: [...added, ...v.levels[level].origin] }, + { type: "translation", translation: [...added, ...translation(v, level, lo)] }, ] })), }; } @@ -54,14 +68,16 @@ export function attrs2(name: string, v: ViewInfo, datasets: Dataset[], extra: Re export const dimensionNames = (v: ViewInfo) => [...(v.added ?? []).map(() => "c"), ...v.axes.map((a) => a.name)]; /** A level's zarr v3 array (its zarr.json), chunked `chunk` (the level's own axes), the - * channel axes the ops add one chunk each; its bytes little endian, then `codecs`. In the - * GeoZarr layout (`map`), the map's y, x band alone, the view's one plane along z (a map - * client reads 2-D bands: with a third axis it would need the group's consolidated - * metadata to know to select it). */ -export function array3(v: ViewInfo, chunk: number[], level: number, { map = false, codecs = [] }: { map?: boolean; codecs?: Codec[] } = {}) { + * channel axes the ops add one chunk each; its bytes little endian, then `codecs`; `shape` + * on the level's own axes (a box of it), the level's by default. In the GeoZarr layout + * (`map`), the map's y, x band alone, the view's one plane along z (a map client reads 2-D + * bands: with a third axis it would need the group's consolidated metadata to know to + * select it). */ +export function array3(v: ViewInfo, chunk: number[], level: number, + { map = false, codecs = [], shape = v.levels[level].shape }: { map?: boolean; codecs?: Codec[]; shape?: number[] } = {}) { const keep = map ? -2 : 0, added = map ? [] : v.added ?? []; return { - zarr_format: 3, node_type: "array", shape: [...added, ...v.levels[level].shape.slice(keep)], data_type: v.out, fill_value: 0, + zarr_format: 3, node_type: "array", shape: [...added, ...shape.slice(keep)], data_type: v.out, fill_value: 0, chunk_grid: { name: "regular", configuration: { chunk_shape: [...added, ...chunk.slice(keep)] } }, chunk_key_encoding: { name: "default", configuration: { separator: "/" } }, codecs: [{ name: "bytes", configuration: { endian: "little" } }, ...codecs], @@ -70,12 +86,13 @@ export function array3(v: ViewInfo, chunk: number[], level: number, { map = fals } /** A level's zarr v2 array (.zarray), chunked `chunk`, the channel axes the ops add one - * chunk each, compressed by `compressor` (numcodecs' config; null: as they are). Its axes' - * names are its .zattrs, `{"_ARRAY_DIMENSIONS": dimensionNames(v)}`. */ -export function array2(v: ViewInfo, chunk: number[], level: number, compressor: Record | null = null) { + * chunk each, compressed by `compressor` (numcodecs' config; null: as they are); `shape` on + * the level's own axes (a box of it), the level's by default. Its axes' names are its + * .zattrs, `{"_ARRAY_DIMENSIONS": dimensionNames(v)}`. */ +export function array2(v: ViewInfo, chunk: number[], level: number, compressor: Record | null = null, shape = v.levels[level].shape) { const added = v.added ?? []; return { - zarr_format: 2, shape: [...added, ...v.levels[level].shape], chunks: [...added, ...chunk], dtype: V2_DTYPE[v.out] ?? v.out, + zarr_format: 2, shape: [...added, ...shape], chunks: [...added, ...chunk], dtype: V2_DTYPE[v.out] ?? v.out, compressor, fill_value: 0, order: "C", filters: null, dimension_separator: "/", }; } diff --git a/web/src/materialize.ts b/web/src/materialize.ts index b16dc8a..74d1ed9 100644 --- a/web/src/materialize.ts +++ b/web/src/materialize.ts @@ -3,11 +3,12 @@ // copy where other tools read it. The store is a directory (one the visitor picks with the // File System Access API's showDirectoryPicker, in Chrome and Edge, or the origin's private // file system) or memory. What is written is what the engine serves for the view -// (layout.ts), one level of it, its chunks compressed. Here are the parts that need no -// engine, so Node's tests run them: the chunks a box covers, the files written, gzip, and -// the loop that has the chunks made and writes them, from any producer (the engine's queue -// on the page, a fake one in the tests). -import { array2, array3, attrs2, chunkKey, dimensionNames, group3 } from "./layout.ts"; // with its extension: Node's runner resolves no other +// (layout.ts), one level of it, or exactly a box of it, its chunks compressed. Here are the +// parts that need no engine, so Node's tests run them: the chunks a box covers, the files +// written, gzip, the array's chunks stitched from the view's for a box that does not start +// on chunk edges, and the loop that has the chunks made and writes them, from any producer +// (the engine's queue on the page, a fake one in the tests). +import { array2, array3, attrs2, chunkKey, dimensionNames, group3, ITEM_BYTES } from "./layout.ts"; // with its extension: Node's runner resolves no other import type { ViewInfo } from "./types.ts"; /** Where materialize writes: files by "/"-separated path relative to the store's root. */ @@ -75,17 +76,31 @@ export function memoryStore(): WritableStore & { files: Map }; } -/** How a materialize is going: chunks done (written or skipped as empty) of `total`, and of - * them `written`, in `bytes` as stored. */ -export interface MaterializeProgress { done: number; total: number; written: number; bytes: number } +/** How a materialize is going: of the written array's `total` chunks, those `done` (written, + * or skipped as empty), and of them those `written`, in `bytes` as stored; of the `needed` + * chunks of the view the box covers, those `computed` so far. Each count only grows. With the + * level's extent the array's chunks are the view's, so `computed` leads `done` by the chunks + * being written; with the box's, an array chunk is done once every chunk of the view it is + * stitched from has been computed, so `done` can trail `computed` by a slab of chunks. */ +export interface MaterializeProgress { done: number; total: number; written: number; bytes: number; computed: number; needed: number } /** `Engine.materialize`'s options. */ export interface MaterializeOptions { /** The view's level to write (default 0, the finest). */ level?: number; /** Voxel box [lo, hi) on that level, in the level's voxel axes (z, y, x...); default the - * whole level. Rounded out to whole chunks, and clipped to the level. */ + * whole level. Clipped to the level, and rounded out to whole voxels (to whole chunks, with + * the level's `extent`). */ lo?: number[]; hi?: number[]; + /** What the written array covers. "level" (the default): the level's whole shape, with the + * level's translation, so the box keeps its place in it; the chunks the box touches are + * written, rounded out to whole chunks. "box": exactly the box, not rounded to chunks: its + * shape hi − lo, its translation the level's moved by lo voxels, cut in the view's chunk + * shape counted from lo. Where the box does not start on the view's chunk edges, each of + * its chunks is stitched from the view's chunks it overlaps (up to 2 along each axis), each + * of those computed once; the chunks at its far edge overhang it and are stored whole, the + * fill value past the box, as zarr stores a chunk at an array's edge. */ + extent?: "level" | "box"; /** "zarr2" (OME-Zarr 0.4: .zgroup, .zattrs, .zarray, "/" as the dimension separator) or * "zarr3" (OME-Zarr 0.5, zarr.json); default "zarr2", which more tools read. */ format?: "zarr3" | "zarr2"; @@ -102,15 +117,17 @@ export interface MaterializeOptions { /** Aborted: no chunk is started after it, chunks waiting for their turn are given up, the * writes under way end, and materialize resolves with `aborted`; chunks written stay. */ signal?: AbortSignal; - /** Told once the metadata is written (done 0), then after each chunk. */ + /** Told once the metadata is written (done 0, computed 0), then after each chunk of the + * view computed, once the array's chunks it completes are written. */ onProgress?: (p: MaterializeProgress) => void; /** Attributes beside the OME multiscales in the group's attributes (the consumer's * provenance); `ome` (zarr v3) or `multiscales` (zarr v2) among them is the view's. */ attributes?: Record; } -/** What a materialize wrote: of the box's `chunks`, those `written` (in `bytes`, as - * stored) and those `skipped` as empty; `aborted` if its signal stopped it first. */ +/** What a materialize wrote: of the written array's `chunks` in the box, those `written` + * (in `bytes`, as stored) and those `skipped` as empty; `aborted` if its signal stopped it + * first. */ export interface Materialized { chunks: number; written: number; skipped: number; bytes: number; aborted: boolean } /** The written level's path in the store: its one dataset, "0", as an OME-Zarr's finest @@ -120,10 +137,10 @@ export const LEVEL_PATH = "0"; * needs none of it to decompress). */ export const GZIP_LEVEL = 6; -/** The chunk grid's indices [lo, hi) covering voxels [lo, hi) of a level of `shape` cut in - * chunks of `chunk`: the box rounded out to whole chunks and clipped to the level; the - * whole level where `lo` or `hi` is left out. A box with no voxel of the level is refused. */ -export function chunkRange(shape: number[], chunk: number[], lo?: number[], hi?: number[]): { lo: number[]; hi: number[] } { +/** Voxels [lo, hi) of a level of `shape`: the box clipped to the level and rounded out to + * whole voxels; the whole level where `lo` or `hi` is left out. A box with no voxel of the + * level is refused. */ +export function voxelBox(shape: number[], lo?: number[], hi?: number[]): { lo: number[]; hi: number[] } { for (const [what, box] of [["lo", lo], ["hi", hi]] as const) { if (box && (box.length !== shape.length || box.some((x) => typeof x !== "number" || Number.isNaN(x)))) { throw new Error(`${what} ${JSON.stringify(box)}: one number for each of the level's ${shape.length} axes`); @@ -134,7 +151,15 @@ export function chunkRange(shape: number[], chunk: number[], lo?: number[], hi?: if (inHi.some((h, a) => h <= inLo[a])) { throw new Error(`the box from [${(lo ?? shape.map(() => 0)).join(", ")}] to [${(hi ?? shape).join(", ")}] holds no voxel of the level, which is ${shape.join("×")}`); } - return { lo: inLo.map((x, a) => Math.floor(x / chunk[a])), hi: inHi.map((x, a) => Math.ceil(x / chunk[a])) }; + return { lo: inLo.map((x) => Math.floor(x)), hi: inHi.map((x) => Math.ceil(x)) }; +} + +/** The chunk grid's indices [lo, hi) covering voxels [lo, hi) of a level of `shape` cut in + * chunks of `chunk`: the box rounded out to whole chunks and clipped to the level; the + * whole level where `lo` or `hi` is left out. A box with no voxel of the level is refused. */ +export function chunkRange(shape: number[], chunk: number[], lo?: number[], hi?: number[]): { lo: number[]; hi: number[] } { + const box = voxelBox(shape, lo, hi); + return { lo: box.lo.map((x, a) => Math.floor(x / chunk[a])), hi: box.hi.map((x, a) => Math.ceil(x / chunk[a])) }; } /** The chunk indices [lo, hi), in C order (the last axis fastest). */ @@ -169,58 +194,169 @@ export function allZero(bytes: Uint8Array): boolean { } /** What a materialize writes of view `name`: which `level`, in which `format` and - * `compression`; the metadata `files` (path and bytes, the group's first), the chunk grid's - * indices [lo, hi) the box covers, `total` chunks, and each chunk's `key` in the store. */ + * `compression`, over which `extent`; the metadata `files` (path and bytes, the group's + * first); the view's chunks the box covers, the chunk grid's indices [lo, hi), `needed` of + * them; and the array written: its voxel 0 the level's voxel `at` (0 for the level's extent, + * the box's lo for the box's), `shape` voxels on the level's axes (the level's own, + * `levelShape`, or the box's), cut in chunks of `chunk` (the view's), those of its grid's + * indices [first, last) written, `total` of them, each chunk the channels the ops add first, + * `channels` values for each voxel, of `item` bytes each; and each chunk's `key` in the + * store. With the level's extent the array's grid is the view's, [first, last) = [lo, hi). */ export interface Region { - name: string; level: number; format: "zarr3" | "zarr2"; compression: "gzip" | null; + name: string; level: number; format: "zarr3" | "zarr2"; compression: "gzip" | null; extent: "level" | "box"; files: [string, Uint8Array][]; - lo: number[]; hi: number[]; total: number; + lo: number[]; hi: number[]; needed: number; + at: number[]; shape: number[]; levelShape: number[]; chunk: number[]; + first: number[]; last: number[]; total: number; + channels: number; item: number; key: (index: number[]) => string; } +const prod = (xs: number[]) => xs.reduce((n, x) => n * x, 1); +/** The count of a grid's indices [lo, hi). */ +const count = (lo: number[], hi: number[]) => prod(lo.map((l, a) => hi[a] - l)); + /** What `options` asks to write of view `name` (served as `v`, chunked `chunk`): the group * listing the one level written as its one dataset, at LEVEL_PATH, with the level's scale and - * translation; the level's array with its whole shape (so the box keeps its place in it), - * the chunks and dtype served, compressed as asked; and the chunks of the box. */ + * translation (moved to the box's first voxel, for the box's extent); the array, the level's + * whole shape (so a box keeps its place in it) or the box's, the chunks and dtype served, + * compressed as asked; and the chunks of the box. */ export function region(name: string, v: ViewInfo, chunk: number[], options: MaterializeOptions = {}): Region { - const { level = 0, format = "zarr2", compression = "gzip", attributes = {} } = options; + const { level = 0, format = "zarr2", compression = "gzip", extent = "level", attributes = {} } = options; if (!Number.isInteger(level) || !v.levels[level]) throw new Error(`${name} has no level ${level}: it has ${v.levels.length}`); if (format !== "zarr2" && format !== "zarr3") throw new Error(`format ${JSON.stringify(format)}: "zarr2" or "zarr3"`); if (compression !== "gzip" && compression !== null) throw new Error(`compression ${JSON.stringify(compression)}: "gzip" or null`); - const { lo, hi } = chunkRange(v.levels[level].shape, chunk, options.lo, options.hi); - const datasets = [{ path: LEVEL_PATH, level }], added = (v.added ?? []).length; + if (extent !== "level" && extent !== "box") throw new Error(`extent ${JSON.stringify(extent)}: "level" or "box"`); + const levelShape = v.levels[level].shape, box = voxelBox(levelShape, options.lo, options.hi); + const { lo, hi } = chunkRange(levelShape, chunk, box.lo, box.hi); // the view's chunks the box needs + const whole = extent === "level", item = ITEM_BYTES[v.out]; + if (!whole && !item) throw new Error(`${name} is ${v.out}, which a box is not cut from`); + const at = whole ? levelShape.map(() => 0) : box.lo, shape = whole ? levelShape : box.hi.map((h, a) => h - box.lo[a]); + const first = whole ? lo : shape.map(() => 0), last = whole ? hi : shape.map((n, a) => Math.ceil(n / chunk[a])); + const datasets = [{ path: LEVEL_PATH, level, ...(whole ? {} : { lo: box.lo }) }], added = v.added ?? []; const files: [string, unknown][] = format === "zarr3" ? [ ["zarr.json", group3(name, v, datasets, attributes)], - [`${LEVEL_PATH}/zarr.json`, array3(v, chunk, level, { codecs: compression ? [{ name: "gzip", configuration: { level: GZIP_LEVEL } }] : [] })], + [`${LEVEL_PATH}/zarr.json`, array3(v, chunk, level, { codecs: compression ? [{ name: "gzip", configuration: { level: GZIP_LEVEL } }] : [], shape })], ] : [ [".zgroup", { zarr_format: 2 }], [".zattrs", attrs2(name, v, datasets, attributes)], - [`${LEVEL_PATH}/.zarray`, array2(v, chunk, level, compression ? { id: "gzip", level: GZIP_LEVEL } : null)], + [`${LEVEL_PATH}/.zarray`, array2(v, chunk, level, compression ? { id: "gzip", level: GZIP_LEVEL } : null, shape)], [`${LEVEL_PATH}/.zattrs`, { _ARRAY_DIMENSIONS: dimensionNames(v) }], ]; const json = (o: unknown) => new TextEncoder().encode(`${JSON.stringify(o, null, 2)}\n`); return { - name, level, format, compression, files: files.map(([path, o]) => [path, json(o)]), - lo, hi, total: lo.reduce((n, l, a) => n * (hi[a] - l), 1), - key: (index) => chunkKey(format, LEVEL_PATH, added, index), + name, level, format, compression, extent, files: files.map(([path, o]) => [path, json(o)]), + lo, hi, needed: count(lo, hi), + at, shape, levelShape, chunk, first, last, total: count(first, last), + channels: prod(added), item: item ?? 0, + key: (index) => chunkKey(format, LEVEL_PATH, added.length, index), + }; +} + +/** The bytes of a chunk of region `r`'s array (or of the view's: they are cut alike). */ +export const chunkBytes = (r: Region) => r.channels * prod(r.chunk) * r.item; + +/** The array's chunk that `add` completes: its index in the array's grid and its bytes, null + * for nothing but the fill value. */ +export type Completed = [index: number[], bytes: Uint8Array | null]; + +/** Region `r`'s array chunks made of the view's chunks, as these come, in any order: `add` + * takes one of the view's chunks (its index in the view's grid and its whole bytes) and + * gives the array's chunks it completes. An array chunk that is one of the view's chunks as + * it is (the same voxels, and its end the level's end where it overhangs: every chunk with + * the level's extent, and those of a box on the view's chunk edges) is that chunk's bytes, + * not copied. Any other is stitched in a buffer of its own, made when the first of the view's + * chunks it overlaps that holds a value not zero comes (none, if they are all zero), each + * one's part copied in, the rest the fill value; it is complete, and forgotten here, once all + * of them have come. The array's chunks begun and not complete are `open`: with the view's + * chunks added in C order, at most a slab of them (one chunk thick along the first axis on + * which the box does not start on a chunk edge, the box's chunks across the axes after it), + * a row and one more, so a row and one where the box starts on chunk edges along z alone, + * one along z and y, and none along every axis; a few more where they come out of order. */ +export function stitcher(r: Region): { add: (index: number[], bytes: Uint8Array) => Completed[]; readonly open: number } { + const C = r.chunk, n = C.length, size = chunkBytes(r); + const open = new Map(); + /** The level's voxels [from, to) of the array's chunk `o` along axis `a`, within the array. */ + const outer = (o: number, a: number) => [r.at[a] + o * C[a], r.at[a] + Math.min((o + 1) * C[a], r.shape[a])]; + /** The view's chunks that the array's chunk `o` is made of. */ + const contributors = (o: number[]) => prod(o.map((x, a) => { + const [from, to] = outer(x, a); + return Math.floor((to - 1) / C[a]) - Math.floor(from / C[a]) + 1; + })); + return { + get open() { return open.size; }, + add(s, bytes) { + // the array's chunks the view's chunk s overlaps, along each axis, within those written + const lo = s.map((x, a) => Math.max(r.first[a], Math.floor((x * C[a] - r.at[a]) / C[a]))); + const hi = s.map((x, a) => Math.min(r.last[a], Math.ceil(((x + 1) * C[a] - r.at[a]) / C[a]))); + const done: Completed[] = []; + let empty: boolean | undefined; // whether s is all zero, once asked + for (const o of chunkIndices(lo, hi)) { + // along each axis: the voxels of s that land in o, from s's voxel src, to o's voxel dst + const src: number[] = [], dst: number[] = [], len: number[] = []; + let same = true; + for (let a = 0; a < n; a++) { + const [from, to] = outer(o[a], a), start = Math.max(s[a] * C[a], from), end = Math.min((s[a] + 1) * C[a], to); + src.push(start - s[a] * C[a]); dst.push(start - from); len.push(end - start); + // the same voxels, and where o overhangs the array, the array ends where the level does + same &&= from === s[a] * C[a] && to === Math.min((s[a] + 1) * C[a], r.levelShape[a]); + } + if (len.some((l) => l <= 0)) continue; // none of s in o's part of the array + if (same) { done.push([o, bytes]); continue; } + if (bytes.byteLength !== size) throw new Error(`${bytes.byteLength} bytes, not a whole chunk's ${size}`); + const key = o.join("."); + let part = open.get(key); + if (!part) open.set(key, part = { bytes: null, left: contributors(o) }); + if (!(empty ??= allZero(bytes))) place(part.bytes ??= new Uint8Array(size), bytes, C, r.channels, r.item, src, dst, len); + if (--part.left === 0) { open.delete(key); done.push([o, part.bytes]); } + } + return done; + }, }; } +/** The voxels `len` long from `src` of chunk bytes `from` copied to `dst` of `to`, both cut in + * chunks of `chunk` with `channels` values of `item` bytes for each voxel, the channels first + * (C order): each run along the last axis copied whole. */ +function place(to: Uint8Array, from: Uint8Array, chunk: number[], channels: number, item: number, src: number[], dst: number[], len: number[]) { + const n = chunk.length, stride = Array(n); + stride[n - 1] = item; + for (let a = n - 2; a >= 0; a--) stride[a] = stride[a + 1] * chunk[a + 1]; + const plane = stride[0] * chunk[0], run = len[n - 1] * item, at = Array(n - 1).fill(0); + for (let c = 0; c < channels; c++) { + for (;;) { + let s = c * plane + src[n - 1] * item, d = c * plane + dst[n - 1] * item; + for (let a = 0; a < n - 1; a++) { s += (src[a] + at[a]) * stride[a]; d += (dst[a] + at[a]) * stride[a]; } + to.set(from.subarray(s, s + run), d); + let a = n - 2; + while (a >= 0 && ++at[a] >= len[a]) { at[a] = 0; a--; } + if (a < 0) break; + } + } +} + /** Chunk `index` of the level as whole bytes, as served (zarr v3's: little endian, C order, - * the channels first); `signal` aborts when it is no longer wanted. */ + * the channels first); `signal` aborts when it is no longer wanted. Asked for once for each + * of the view's chunks a region needs. */ export type Produce = (index: number[], signal: AbortSignal) => Promise; const why = (e: unknown) => String((e as Error)?.message ?? e); -/** Write region `r` into `store`: its metadata, then its chunks, `concurrency` of them asked - * of `produce` at once (one by default), each as it comes, gzipped if `r` says so, those of - * nothing but zeros skipped (`skipEmpty`, the default). Resolves with what was written, or - * on `signal`'s abort with `aborted` (no chunk started after it, the ones asked for given up - * through `produce`'s signal and not waited for, the writes under way ended). A chunk that - * fails to be made or written rejects it, naming the chunk and why, once the writes under - * way have ended; the rest are given up as on an abort. */ +/** Write region `r` into `store`: its metadata, then its chunks. The view's chunks the box + * covers are asked of `produce` once each, in C order (the last axis fastest), + * `concurrency` at once (one by default); each, as it comes, completes the array's chunks it + * is the last of (`stitcher`: with the level's extent, or a box on the view's chunk edges, + * just itself), and those are written, gzipped if `r` says so, those of nothing but zeros + * skipped (`skipEmpty`, the default). Only the array's chunks begun and not yet complete are + * held meanwhile (about a slab of them at most, `stitcher`), with the view's chunks on their + * way. + * Resolves with what was written, or on `signal`'s abort with `aborted` (no chunk asked for + * or written after it, the ones asked for given up through `produce`'s signal and not waited + * for, the writes under way ended). A chunk that fails to be made or written rejects it, + * naming the chunk and why, once the writes under way have ended; the rest are given up as on + * an abort. */ export async function writeChunks(store: WritableStore, r: Region, produce: Produce, options: Pick = {}): Promise { const { skipEmpty = true, concurrency = 1, signal, onProgress } = options; @@ -237,38 +373,46 @@ export async function writeChunks(store: WritableStore, r: Region, produce: Prod stopped.catch(() => {}); let failure: Error | null = null; const fail = (e: Error) => { failure ??= e; stop.abort(e); }; + let computed = 0; const report = () => { - try { onProgress?.({ done: got.written + got.skipped, total: r.total, written: got.written, bytes: got.bytes }); } + try { onProgress?.({ done: got.written + got.skipped, total: r.total, written: got.written, bytes: got.bytes, computed, needed: r.needed }); } catch (e) { console.error(`materialize ${r.name}: onProgress threw: ${why(e)}`); } // the page's own trouble: the writing goes on }; report(); - const indices = chunkIndices(r.lo, r.hi); + // a chunk of the view named as the array's chunk it is, with the level's extent; as the engine names it, with the box's + const named = (index: number[]) => (r.extent === "level" ? `chunk ${r.key(index)} of ${r.name}` : `chunk ${r.name}/${r.level}/${index.join(".")}`); + const stitched = stitcher(r), indices = chunkIndices(r.lo, r.hi); const lane = async () => { for (let next = indices.next(); !next.done && !stop.signal.aborted; next = indices.next()) { - const index = next.value, key = r.key(index); - let body: ArrayBuffer; + const index = next.value; + let body: ArrayBuffer, complete: Completed[]; try { body = await Promise.race([produce(index, stop.signal), stopped]); // given up at once when stopped } catch (e) { - if (!stop.signal.aborted) fail(new Error(`chunk ${key} of ${r.name}: ${why(e)}`)); + if (!stop.signal.aborted) fail(new Error(`${named(index)}: ${why(e)}`)); return; } if (stop.signal.aborted) return; // made as it stopped: not written - const bytes = new Uint8Array(body); - if (skipEmpty && allZero(bytes)) { got.skipped++; report(); continue; } - try { - const stored = r.compression === "gzip" ? await gzip(bytes) : bytes; - await store.put(key, stored); - got.written++; got.bytes += stored.byteLength; - } catch (e) { - fail(new Error(`writing chunk ${key} of ${r.name}: ${why(e)}`)); - return; + computed++; + try { complete = stitched.add(index, new Uint8Array(body)); } catch (e) { fail(new Error(`${named(index)}: ${why(e)}`)); return; } + for (const [o, bytes] of complete) { + if (stop.signal.aborted) return; // no write starts once stopped + const key = r.key(o); + if (skipEmpty && (!bytes || allZero(bytes))) { got.skipped++; continue; } + try { + const raw = bytes ?? new Uint8Array(chunkBytes(r)), stored = r.compression === "gzip" ? await gzip(raw) : raw; + await store.put(key, stored); + got.written++; got.bytes += stored.byteLength; + } catch (e) { + fail(new Error(`writing chunk ${key} of ${r.name}: ${why(e)}`)); + return; + } } report(); } }; try { - await Promise.all(Array.from({ length: Math.min(concurrency, r.total) }, lane)); + await Promise.all(Array.from({ length: Math.min(concurrency, r.needed) }, lane)); } finally { signal?.removeEventListener("abort", onAbort); } diff --git a/web/test/materialize.test.ts b/web/test/materialize.test.ts index 08ff27f..bde35c7 100644 --- a/web/test/materialize.test.ts +++ b/web/test/materialize.test.ts @@ -2,13 +2,15 @@ // that need no engine): the chunks a box covers, rounded out and clipped; the metadata both // formats write, the layout the engine serves (src/layout.ts); gzip; the loop over a fake // producer, with empty chunks skipped, progress, abort and failure; the stores; and the -// result read back by zarrita, a zarr implementation of its own, as a reader would. Run by -// `npm test`. +// result read back by zarrita, a zarr implementation of its own, as a reader would. Then the +// same for a box written as exactly the box (extent "box"): its metadata, its chunks +// stitched from the view's (each of those made once), its edge chunks, and the slab of +// chunks held at most. Run by `npm test`. import { test } from "node:test"; import assert from "node:assert/strict"; import * as zarr from "zarrita"; import { array3, chunkKey, group3 } from "../src/layout.ts"; -import { allZero, chunkIndices, chunkRange, directoryStore, gzip, memoryStore, region, writeChunks, type MaterializeProgress, type Produce, type WritableStore } from "../src/materialize.ts"; +import { allZero, chunkIndices, chunkRange, directoryStore, gzip, memoryStore, region, stitcher, voxelBox, writeChunks, type MaterializeOptions, type MaterializeProgress, type Produce, type WritableStore } from "../src/materialize.ts"; import type { ViewInfo } from "../src/types.ts"; const nm = (name: string) => ({ name, unit: "nm" }); @@ -172,8 +174,8 @@ for (const format of ["zarr2", "zarr3"] as const) { assert.deepEqual(got, { chunks: 8, written: 7, skipped: 1, bytes, aborted: false }); assert.deepEqual(await gunzip(store.files.get(at("1/1/1"))!), new Uint8Array(fakeChunk([1, 1, 1]))); // progress: once with the metadata written, then after each chunk, done = written + skipped - assert.deepEqual(seen[0], { done: 0, total: 8, written: 0, bytes: 0 }); - assert.deepEqual(seen.at(-1), { done: 8, total: 8, written: 7, bytes }); + assert.deepEqual(seen[0], { done: 0, total: 8, written: 0, bytes: 0, computed: 0, needed: 8 }); + assert.deepEqual(seen.at(-1), { done: 8, total: 8, written: 7, bytes, computed: 8, needed: 8 }); const steps = seen.slice(1).map((p, i) => [p.done - seen[i].done, p.written - seen[i].written, p.bytes - seen[i].bytes]); assert.deepEqual(steps.map(([d]) => d), [1, 1, 1, 1, 1, 1, 1, 1]); assert.deepEqual(steps.filter(([, w, b]) => w === 0 && b === 0).length, 1); // the one skipped @@ -311,3 +313,287 @@ test("a directory as a store: the root and subdirectories made once each, each f assert.ok(calls.includes("aborted bad")); assert.equal(failing.tree.files.size, 0); }); + +// ------------------------------------------------------------- a box as exactly the box + +/** Voxel (z, y, x) of channel c of level 0 as the box tests' producer makes it: never zero + * inside the level, and no two alike. */ +const voxel = (z: number, y: number, x: number, c = 0) => 1 + 10000 * c + 1000 * z + 20 * y + x; +/** `view` with a level 0 of 9 × 10 × 11 voxels: 5 × 3 × 3 chunks of 2 × 4 × 4, the last + * along each axis partial. */ +const big: ViewInfo = { ...view, levels: [{ shape: [9, 10, 11], voxel: [8, 8, 8], origin: [0, 0, 0] }, view.levels[1]] }; +/** `big` with two channels its ops add, as float32. */ +const bigChannels: ViewInfo = { ...big, out: "float32", added: [2] }; + +/** Level 0 of `v` as the engine gives its chunks: whole, zero past the level's edge, the + * channels the ops add first; the chunks of `empty` all zero; each after `delay(index)` + * ticks. `calls` counts the asks for each chunk, `asked` lists them in order. */ +function served(v: ViewInfo, { empty = [] as string[], delay = (_: number[]) => 0 } = {}) { + const [Z, Y, X] = v.levels[0].shape, K = (v.added ?? []).reduce((n, c) => n * c, 1), n = chunk[0] * chunk[1] * chunk[2]; + const make = (index: number[]): ArrayBuffer => { + const out = v.out === "float32" ? new Float32Array(K * n) : new Uint16Array(K * n); + if (empty.includes(index.join("."))) return out.buffer; + for (let c = 0; c < K; c++) for (let z = 0; z < chunk[0]; z++) for (let y = 0; y < chunk[1]; y++) for (let x = 0; x < chunk[2]; x++) { + const g = [index[0] * chunk[0] + z, index[1] * chunk[1] + y, index[2] * chunk[2] + x]; + if (g[0] < Z && g[1] < Y && g[2] < X) out[c * n + (z * chunk[1] + y) * chunk[2] + x] = voxel(g[0], g[1], g[2], c); + } + return out.buffer; + }; + const calls = new Map(), asked: string[] = []; + const produce: Produce = async (index) => { + const k = index.join("."); + calls.set(k, (calls.get(k) ?? 0) + 1); + asked.push(k); + for (let i = delay(index); i > 0; i--) await tick(); + return make(index); + }; + return { produce, make, calls, asked }; +} + +/** The array at "0" of `store` as zarrita reads it. */ +async function readArray(store: { files: Map }, format: "zarr2" | "zarr3") { + const at = zarr.root(readable(store)).resolve("0"); + const arr = format === "zarr3" ? await zarr.open.v3(at, { kind: "array" }) : await zarr.open.v2(at, { kind: "array" }); + return { shape: arr.shape, dtype: arr.dtype, data: (await zarr.get(arr)).data as Uint16Array | Float32Array }; +} + +/** That `data`, of `shape` (its channels first), is the view's voxels from voxel `lo` of the + * level; zero where `zero` says (a chunk of the view left empty). */ +function assertBox(data: ArrayLike, shape: number[], lo: number[], zero: (g: number[]) => boolean = () => false) { + const [Z, Y, X] = shape.slice(-3), K = shape.length > 3 ? shape[0] : 1; + for (let c = 0; c < K; c++) for (let z = 0; z < Z; z++) for (let y = 0; y < Y; y++) for (let x = 0; x < X; x++) { + const g = [lo[0] + z, lo[1] + y, lo[2] + x]; + assert.equal(data[((c * Z + z) * Y + y) * X + x], zero(g) ? 0 : voxel(g[0], g[1], g[2], c), `channel ${c}, voxel ${g}`); + } +} + +/** The C-order cells of the grid [lo, hi), as "z.y.x". */ +const cells = (lo: number[], hi: number[]) => [...chunkIndices(lo, hi)].map((i) => i.join(".")); + +test("a box's array is exactly the box: its shape hi − lo, the view's chunks counted from lo, the level's translation moved by lo voxels", () => { + // level 1: 3 × 3 × 4 voxels of 16 nm, voxel 0's centre at 4 nm + const two = region("mito", view, chunk, { level: 1, extent: "box", lo: [1, 0, 2], hi: [3, 2, 4] }); + const [, attrs, zarray] = two.files.map(([, b]) => json(b)); + assert.deepEqual(attrs.multiscales[0].datasets, [{ path: "0", coordinateTransformations: [ + { type: "scale", scale: [16, 16, 16] }, { type: "translation", translation: [4 + 16, 4, 4 + 2 * 16] }] }]); + assert.deepEqual([zarray.shape, zarray.chunks], [[2, 2, 2], [2, 4, 4]]); + // the view's chunks it needs (z 0 and 1), the array's one chunk + assert.deepEqual([two.lo, two.hi, two.needed, two.first, two.last, two.total, two.key([0, 0, 0])], [[0, 0, 0], [2, 1, 1], 2, [0, 0, 0], [1, 1, 1], 1, "0/0/0/0"]); + const three = region("mito", view, chunk, { format: "zarr3", extent: "box", lo: [1, 3, 2], hi: [4, 6, 7] }); + const [group, array] = three.files.map(([, b]) => json(b)); + assert.deepEqual(group.attributes.ome.multiscales[0].datasets[0].coordinateTransformations[1], { type: "translation", translation: [8, 24, 16] }); + assert.deepEqual([array.shape, array.chunk_grid.configuration.chunk_shape], [[3, 3, 5], [2, 4, 4]]); + assert.deepEqual([three.needed, three.total, three.key([1, 0, 1])], [2 * 2 * 2, 2 * 1 * 2, "0/c/1/0/1"]); + // clipped to the level and rounded out to whole voxels, not to chunks + assert.deepEqual(voxelBox([5, 6, 7], [-3, 0.5, 1.5], [2.5, 100, 3.2]), { lo: [0, 0, 1], hi: [3, 6, 4] }); + assert.deepEqual(region("v", view, chunk, { extent: "box", lo: [-3, 0.5, 1.5], hi: [2.5, 100, 3.2] }).shape, [3, 6, 3]); + assert.throws(() => region("v", view, chunk, { extent: "box", lo: [5, 0, 0] }), /holds no voxel of the level, which is 5×6×7/); + assert.throws(() => region("v", view, chunk, { extent: "slab" as "box" }), /extent "slab": "level" or "box"/); + // the channels the ops add: whole, first, untranslated + const grad = region("grad", channels, chunk, { extent: "box", lo: [1, 3, 2], hi: [4, 6, 7] }); + const [, gAttrs, gArray] = grad.files.map(([, b]) => json(b)); + assert.deepEqual([gArray.shape, gArray.chunks], [[2, 3, 3, 5], [2, 2, 4, 4]]); + assert.deepEqual(gAttrs.multiscales[0].datasets[0].coordinateTransformations[1].translation, [0, 8, 24, 16]); + assert.equal(grad.key([1, 0, 1]), "0/0/1/0/1"); + // the level's own extent is as before: the level's whole shape and translation + const level = region("mito", view, chunk, { level: 1, lo: [1, 0, 2], hi: [3, 2, 4] }); + assert.deepEqual(json(level.files[2][1]).shape, [3, 3, 4]); + assert.deepEqual(json(level.files[1][1]).multiscales[0].datasets[0].coordinateTransformations[1].translation, [4, 4, 4]); +}); + +for (const format of ["zarr2", "zarr3"] as const) { + test(`${format}: a box across chunk edges on every axis, each of its chunks stitched from up to 8 of the view's, each of those made once, read back by zarrita as the view over the box`, async () => { + const lo = [1, 3, 2], hi = [8, 9, 11]; // 7 × 6 × 9: 4 × 2 × 3 chunks, of the view's 4 × 3 × 3; to the level's end along x + const fake = served(big, { delay: (i) => (i[0] + 2 * i[1] + i[2]) % 3 }); // made out of turn + const store = memoryStore(), seen: MaterializeProgress[] = []; + const got = await writeChunks(store, region("v", big, chunk, { format, extent: "box", lo, hi }), fake.produce, + { concurrency: 3, onProgress: (p) => seen.push(p) }); + // each of the view's chunks the box needs asked for once, and no other + assert.deepEqual([...fake.calls.keys()].sort(), cells([0, 0, 0], [4, 3, 3]).sort()); + assert.ok([...fake.calls.values()].every((n) => n === 1)); + const keys = files(store).filter((f) => !f.endsWith("json") && !f.includes(".z")); + assert.deepEqual(keys, cells([0, 0, 0], [4, 2, 3]).map((i) => (format === "zarr3" ? "0/c/" : "0/") + i.replaceAll(".", "/")).sort()); + const bytes = keys.reduce((n, k) => n + store.files.get(k)!.length, 0); + assert.deepEqual(got, { chunks: 24, written: 24, skipped: 0, bytes, aborted: false }); + const arr = await readArray(store, format); + assert.deepEqual([arr.shape, arr.dtype], [[7, 6, 9], "uint16"]); + assertBox(arr.data, arr.shape, lo); + // the far corner's chunk overhangs the box (1 of 2 along z, 2 of 4 along y, 1 of 4 along x), + // stored whole, as zarr stores an edge chunk, zero past the box + const corner = new Uint16Array((await gunzip(store.files.get(format === "zarr3" ? "0/c/3/1/2" : "0/3/1/2")!)).buffer); + assert.equal(corner.length, chunk[0] * chunk[1] * chunk[2]); + corner.forEach((v, i) => { + const [z, y, x] = [Math.floor(i / 16), Math.floor(i / 4) % 4, i % 4]; + assert.equal(v, z < 1 && y < 2 && x < 1 ? voxel(lo[0] + 6 + z, lo[1] + 4 + y, lo[2] + 8 + x) : 0, `corner ${z}, ${y}, ${x}`); + }); + // progress: the array's chunks done, the view's computed, each only growing + assert.deepEqual(seen[0], { done: 0, total: 24, written: 0, bytes: 0, computed: 0, needed: 36 }); + assert.deepEqual(seen.at(-1), { done: 24, total: 24, written: 24, bytes, computed: 36, needed: 36 }); + assert.equal(seen.length, 1 + 36); // once more for each of the view's chunks + for (const [i, p] of seen.slice(1).entries()) { + for (const k of ["done", "written", "bytes", "computed"] as const) assert.ok(p[k] >= seen[i][k], `${k} went back at report ${i + 1}`); + assert.ok(p.computed >= i + 1); // another lane's chunk may be counted before it reports + } + }); +} + +test("a box on the view's chunk edges: its chunks are the view's as they are; one ending inside a chunk crops it", async () => { + // from chunk (1, 1, 1) to the level's end: 7 × 6 × 7, the view's chunks (1..4, 1..2, 1..2) + const fake = served(big), store = memoryStore(); + const r = region("v", big, chunk, { extent: "box", compression: null, lo: [2, 4, 4] }); + assert.deepEqual([r.shape, r.lo, r.hi, r.total, r.needed], [[7, 6, 7], [1, 1, 1], [5, 3, 3], 16, 16]); + assert.deepEqual(await writeChunks(store, r, fake.produce), { chunks: 16, written: 16, skipped: 0, bytes: 16 * 64, aborted: false }); + for (const o of chunkIndices([0, 0, 0], [4, 2, 2])) { + assert.deepEqual(store.files.get(r.key(o)), new Uint8Array(fake.make(o.map((i) => i + 1))), `chunk ${o}`); // byte for byte + } + assertBox((await readArray(store, "zarr2")).data, [7, 6, 7], [2, 4, 4]); + const s = stitcher(r); + for (const i of chunkIndices(r.lo, r.hi)) { + assert.equal(s.add(i, new Uint8Array(fake.make(i))).length, 1); + assert.equal(s.open, 0); // none held: each passes through + } + // on chunk edges at its start, but its end inside chunks short of the level's: 3 × 3 × 2, cropped + const short = memoryStore(), cropped = region("v", big, chunk, { extent: "box", compression: null, lo: [2, 4, 4], hi: [5, 7, 6] }); + await writeChunks(short, cropped, served(big).produce); + const second = new Uint16Array(short.files.get("0/1/0/0")!.slice().buffer); // z 4 of 4..5, y 4..6 of 4..7, x 4..5 of 4..7 + second.forEach((v, i) => { + const [z, y, x] = [Math.floor(i / 16), Math.floor(i / 4) % 4, i % 4]; + assert.equal(v, z < 1 && y < 3 && x < 2 ? voxel(4 + z, 4 + y, 4 + x) : 0, `voxel ${z}, ${y}, ${x}: the view's past the box not kept`); + }); + assertBox((await readArray(short, "zarr2")).data, [3, 3, 2], [2, 4, 4]); +}); + +test("a box smaller than a chunk: inside one of the view's, and across the corner of eight", async () => { + const inside = served(view), one = memoryStore(); + const got = await writeChunks(one, region("v", view, chunk, { extent: "box", compression: null, lo: [1, 1, 1], hi: [2, 3, 2] }), inside.produce); + assert.deepEqual([got, [...inside.calls.keys()]], [{ chunks: 1, written: 1, skipped: 0, bytes: 64, aborted: false }, ["0.0.0"]]); + const stored = new Uint16Array(one.files.get("0/0/0/0")!.slice().buffer); // whole, the two voxels of the box first + assert.deepEqual([stored.length, stored[0], stored[4], stored.filter(Boolean).length], [32, voxel(1, 1, 1), voxel(1, 2, 1), 2]); + assertBox((await readArray(one, "zarr2")).data, [1, 2, 1], [1, 1, 1]); + const corner = served(view), eight = memoryStore(); + await writeChunks(eight, region("v", view, chunk, { format: "zarr3", extent: "box", lo: [1, 3, 3], hi: [3, 5, 5] }), corner.produce, { concurrency: 3 }); + assert.deepEqual([...corner.calls.keys()].sort(), cells([0, 0, 0], [2, 2, 2])); + assert.ok([...corner.calls.values()].every((n) => n === 1)); + const arr = await readArray(eight, "zarr3"); + assert.deepEqual(arr.shape, [2, 2, 2]); + assertBox(arr.data, arr.shape, [1, 3, 3]); +}); + +for (const format of ["zarr2", "zarr3"] as const) { + test(`${format}: a box of the whole level is the level's extent, file for file and byte for byte`, async () => { + for (const box of [{}, { lo: [0, 0, 0], hi: [5, 6, 7] }, { lo: [-1, -1, -1], hi: [9, 9, 9] }] as MaterializeOptions[]) { + const level = memoryStore(), whole = memoryStore(); + await writeChunks(level, region("v", view, chunk, { format, ...box }), producer(["1.0.1"])); + await writeChunks(whole, region("v", view, chunk, { format, extent: "box", ...box }), producer(["1.0.1"])); + assert.deepEqual(whole.files, level.files); + } + }); +} + +test("skipEmpty: a box's chunk is left out when the view's chunks it is made of are all empty; one made of some is written, zero there", async () => { + // 4 × 8 × 8 from (1, 1, 1): 2 × 2 × 2 chunks, the first made of the view's (0..1, 0..1, 0..1), all empty + const empty = cells([0, 0, 0], [2, 2, 2]), zero = (g: number[]) => g[0] < 4 && g[1] < 8 && g[2] < 8; + const opts: MaterializeOptions = { extent: "box", lo: [1, 1, 1], hi: [5, 9, 9] }; + const store = memoryStore(); + assert.deepEqual(await writeChunks(store, region("v", big, chunk, opts), served(big, { empty }).produce), + { chunks: 8, written: 7, skipped: 1, bytes: [...store.files].filter(([k]) => /^0\/\d/.test(k)).reduce((n, [, b]) => n + b.length, 0), aborted: false }); + assert.ok(!store.files.has("0/0/0/0") && store.files.has("0/0/0/1")); + assertBox((await readArray(store, "zarr2")).data, [4, 8, 8], [1, 1, 1], zero); + // kept when asked: written whole, all zero + const all = memoryStore(); + const got = await writeChunks(all, region("v", big, chunk, { ...opts, compression: null }), served(big, { empty }).produce, { skipEmpty: false }); + assert.deepEqual([got.written, got.skipped, all.files.get("0/0/0/0")], [8, 0, new Uint8Array(64)]); + assertBox((await readArray(all, "zarr2")).data, [4, 8, 8], [1, 1, 1], zero); +}); + +test("a box's chunk is written once the last of the view's chunks it is made of comes; a slab of them held at most", async () => { + // from (1, 1, 1) to the level's end: off the chunk edges on every axis, 8 × 9 × 10 in 4 × 3 × 3 chunks + const r = region("v", big, chunk, { extent: "box", compression: null, lo: [1, 1, 1] }), events: string[] = []; + const fake = served(big), produce: Produce = async (i, signal) => { events.push(`make ${i.join(".")}`); return fake.produce(i, signal); }; + const store: WritableStore = { async put(path) { if (/^0\/\d/.test(path)) events.push(`write ${path}`); } }; + await writeChunks(store, r, produce); + const writes = events.filter((e) => e.startsWith("write")); + assert.equal(writes.length, r.total); + for (const [i, e] of events.entries()) { + if (!e.startsWith("write")) continue; + // the last of the view's chunks it is made of, in C order: the one each axis's last part comes from + const o = e.slice("write 0/".length).split("/").map(Number); + const last = o.map((x, a) => Math.floor((r.at[a] + Math.min((x + 1) * chunk[a], r.shape[a]) - 1) / chunk[a])); + assert.equal(events.slice(0, i).filter((x) => x.startsWith("make")).at(-1), `make ${last.join(".")}`, `${e} as soon as it is complete`); + } + // held: off chunk edges along z, a slab of the box's chunks (3 × 3), a row (3) and one more at most + const slab = (lo: number[]) => { + const s = stitcher(region("v", big, chunk, { extent: "box", compression: null, lo })), r2 = region("v", big, chunk, { extent: "box", lo }); + let most = 0; + for (const i of chunkIndices(r2.lo, r2.hi)) { s.add(i, new Uint8Array(fake.make(i))); most = Math.max(most, s.open); } + assert.equal(s.open, 0); // all complete at the end + return most; + }; + assert.equal(slab([1, 1, 1]), 3 * 3 + 3 + 1); + assert.equal(slab([2, 1, 1]), 3 + 1); // on edges along z: a row and one + assert.equal(slab([2, 4, 1]), 1); // along z and y: one + assert.equal(slab([2, 4, 4]), 0); + // the view's chunks all empty: nothing held in memory but the counts (no buffer made) + const r3 = region("v", big, chunk, { extent: "box", lo: [1, 1, 1] }), s3 = stitcher(r3); + const done = [...chunkIndices(r3.lo, r3.hi)].flatMap((i) => s3.add(i, new Uint8Array(64))); + assert.deepEqual([done.length, done.every(([, b]) => b === null)], [r3.total, true]); +}); + +test("a box's abort: no chunk asked for or written after it, those asked for given up; the chunks written stay", { timeout: 5000 }, async () => { + const store = memoryStore(), stop = new AbortController(), asked: string[] = [], given: AbortSignal[] = [], fake = served(big); + let atAbort = -1; + // the first 16 of the view's chunks come at once, the rest never; the box's first chunk is complete with the 14th + const produce: Produce = (index, signal) => { + asked.push(index.join(".")); + given.push(signal); + return asked.length <= 16 ? Promise.resolve(fake.make(index)) : new Promise(() => {}); + }; + const got = await writeChunks(store, region("v", big, chunk, { extent: "box", lo: [1, 1, 1] }), produce, { + concurrency: 2, signal: stop.signal, onProgress: (p) => { if (p.written >= 1 && atAbort < 0) { stop.abort(); atAbort = asked.length; } }, + }); + assert.deepEqual([got.chunks, got.aborted], [36, true]); + assert.ok(got.written >= 1); + assert.equal(asked.length, atAbort); // none asked for after the abort + assert.ok(given.every((s) => s.aborted)); + assert.ok(store.files.has("0/0/0/0")); + assert.equal(files(store).filter((f) => /^0\/\d/.test(f)).length, got.written); +}); + +test("a box's chunk that fails rejects naming it as the engine does, once the writes under way have ended; no write starts after", async () => { + // 4 × 5 × 6 from (1, 1, 1) of view: the view's chunk (1, 1, 1) completes four of its chunks, (2, 0, 0) fails + let slowWritten = false; + const memory = memoryStore(); + const store: WritableStore = { + async put(path, bytes) { + if (path === "0/0/0/0") { await new Promise((r) => setTimeout(r, 30)); slowWritten = true; } + return memory.put(path, bytes); + }, + }; + const fake = served(view); + const produce: Produce = async (index, signal) => { + if (index.join(".") === "2.0.0") { await tick(); throw new Error("the model ran out of GPU memory"); } + return fake.produce(index, signal); + }; + await assert.rejects(writeChunks(store, region("v", view, chunk, { extent: "box", lo: [1, 1, 1] }), produce, { concurrency: 2 }), (e: Error) => { + assert.equal(e.message, "chunk v/0/2.0.0: the model ran out of GPU memory"); + assert.ok(slowWritten, "the chunk being written when the other failed was written first"); + return true; + }); + assert.deepEqual(files(memory).filter((f) => /^0\/\d/.test(f)), ["0/0/0/0"]); // not the three it completed with it + // a chunk of the view that is not a whole chunk cannot be cut + await assert.rejects(writeChunks(memoryStore(), region("v", view, chunk, { extent: "box", lo: [1, 1, 1] }), async () => new ArrayBuffer(10)), + /^Error: chunk v\/0\/0\.0\.0: 10 bytes, not a whole chunk's 64$/); +}); + +for (const format of ["zarr2", "zarr3"] as const) { + test(`${format}: a box of a view with channel axes its ops add: the channels whole and first, each stitched alike`, async () => { + const lo = [1, 3, 2], hi = [8, 9, 11], fake = served(bigChannels), store = memoryStore(); + const got = await writeChunks(store, region("grad", bigChannels, chunk, { format, extent: "box", lo, hi }), fake.produce, { concurrency: 2 }); + assert.deepEqual([got.chunks, got.written], [24, 24]); + assert.ok([...fake.calls.values()].every((n) => n === 1) && fake.calls.size === 36); + assert.ok(store.files.has(format === "zarr3" ? "0/c/0/3/1/2" : "0/0/3/1/2")); + const arr = await readArray(store, format); + assert.deepEqual([arr.shape, arr.dtype], [[2, 7, 6, 9], "float32"]); + assertBox(arr.data, arr.shape, lo); + }); +} From 2217e2bec7f135ceb38b02a1f60d299459972da7 Mon Sep 17 00:00:00 2001 From: David Ackerman Date: Fri, 16 Oct 2026 15:39:18 -0400 Subject: [PATCH 2/2] =?UTF-8?q?Engine.materialize's=20box=20is=20the=20def?= =?UTF-8?q?ault,=20walked=20to=20hold=20the=20thinnest=20slab,=20as=20the?= =?UTF-8?q?=20chunkmirage=20session=20asked=20in=20its=20review=20of=20431?= =?UTF-8?q?499b:=20extent=20defaults=20to=20"box",=20since=20a=20page=20th?= =?UTF-8?q?at=20passes=20a=20box=20expects=20an=20array=20of=20that=20box,?= =?UTF-8?q?=20and=20"level"=20stays=20for=20several=20boxes=20written=20in?= =?UTF-8?q?to=20one=20store=20(a=20box=20of=20the=20whole=20level=20is=20t?= =?UTF-8?q?he=20same=20store,=20byte=20for=20byte,=20either=20way);=20the?= =?UTF-8?q?=20served=20chunks=20are=20walked=20axis=20by=20axis=20(walkOrd?= =?UTF-8?q?er,=20chunkIndices=20taking=20an=20order=20of=20the=20axes),=20?= =?UTF-8?q?the=20axes=20the=20box=20does=20not=20split=20(splitAxes:=20lo?= =?UTF-8?q?=20on=20the=20chunk=20edges=20there,=20or=20the=20box=20within?= =?UTF-8?q?=20one=20chunk)=20outermost=20in=20C=20order,=20then=20the=20sp?= =?UTF-8?q?lit=20ones,=20the=20one=20with=20the=20most=20of=20the=20box's?= =?UTF-8?q?=20chunks=20first,=20so=20the=20box's=20chunks=20held=20at=20on?= =?UTF-8?q?ce=20are=20a=20slab=20across=20the=20other=20split=20axes,=20a?= =?UTF-8?q?=20row=20and=20one=20more=20(n2=20=C3=97=20n3=20+=20n3=20+=201?= =?UTF-8?q?=20with=20n1=20=E2=89=A5=20n2=20=E2=89=A5=20n3=20chunks=20along?= =?UTF-8?q?=20three=20split=20axes,=20n2=20+=201=20along=20two,=201=20alon?= =?UTF-8?q?g=20one,=20none=20along=20none),=20the=20fewest=20of=20any=20or?= =?UTF-8?q?der=20of=20the=20axes:=20a=20wide,=20shallow=20box=20in=20C=20o?= =?UTF-8?q?rder=20held=20a=20slab=20across=20its=20two=20wide=20axes=20(a?= =?UTF-8?q?=20512=20=C3=97=202048=20=C3=97=202048=20box=20of=203-channel?= =?UTF-8?q?=20float32=2064=C2=B3=20chunks,=201057=20chunks,=203.1=20GB,=20?= =?UTF-8?q?now=20265,=20795=20MiB;=20on=20a=20chunk=20edge=20along=20z,=20?= =?UTF-8?q?33);=20the=20level's=20extent=20splits=20none,=20so=20it=20is?= =?UTF-8?q?=20walked=20in=20C=20order=20as=20before;=20the=20bytes=20of=20?= =?UTF-8?q?a=20value=20come=20from=20asread's=20TYPED,=20not=20a=20table?= =?UTF-8?q?=20of=20their=20own;=20the=20Python=20side=20says=20it=20writes?= =?UTF-8?q?=20the=20level's=20extent=20only,=20for=20now=20(formats,=20cli?= =?UTF-8?q?);=20the=20tests=20ask=20for=20the=20level's=20extent=20by=20na?= =?UTF-8?q?me=20where=20they=20are=20of=20it,=20and=20check=20the=20defaul?= =?UTF-8?q?t,=20the=20order=20of=20the=20axes,=20the=20asks=20in=20it=20an?= =?UTF-8?q?d=20the=20chunks=20held=20for=20each=20alignment=20and=20order;?= =?UTF-8?q?=20formats,=20design,=20caching=20and=20the=20CHANGELOG=20say?= =?UTF-8?q?=20so?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- CHANGELOG.md | 17 ++++--- docs/concepts/caching.md | 6 +-- docs/concepts/formats.md | 50 +++++++++++++-------- docs/design.md | 18 +++++--- docs/reference/cli.md | 2 +- web/src/engine.ts | 8 ++-- web/src/layout.ts | 4 -- web/src/materialize.ts | 74 ++++++++++++++++++++---------- web/test/materialize.test.ts | 87 +++++++++++++++++++++++++----------- 9 files changed, 173 insertions(+), 93 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 5ad3125..4e11226 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -153,14 +153,17 @@ The first release, an alpha: the plugin API can still change before 0.1.0. - `Engine.materialize(view, store, options)` writes a box of a served view (or a whole level) into a store as an OME-Zarr of its own, zarr v2 with OME 0.4 or zarr v3 with OME 0.5, gzipped, its empty chunks left out, the metadata the engine serves for the view - (`web/src/layout.ts`, shared with the routes) with the one level as its one dataset: the + (`web/src/layout.ts`, shared with the routes) with the one level as its one dataset: + exactly the box (`extent: "box"`, the default), `hi − lo` voxels with the level's + translation moved by `lo` voxels, its chunks the served chunk shape counted from `lo`, + each stitched from the served chunks it overlaps (each computed once, walked with the + split axis of the most chunks outermost, so a slab across the other split axes is held + at a time), those at its far edge stored whole with the fill value past the box; or the level's whole shape, the box's chunks in it rounded out to whole chunks (`extent: - "level"`, the default), or exactly the box (`extent: "box"`), `hi − lo` voxels with the - level's translation moved by `lo` voxels, its chunks the served chunk shape counted from - `lo`, each stitched from the served chunks it overlaps (each computed once, about a slab - of the box's chunks held at a time), those at its far edge stored whole with the fill - value past the box; progress counts the array's chunks done and the served chunks - computed (`computed`, `needed`); + "level"`, for several boxes in one store); a box of the whole level is the same store + either way; progress counts the array's chunks done and the served chunks computed + (`computed`, `needed`); the Python side (`chunkmirage materialize`) writes the level's + extent only, for now; `directoryStore()` writes into a folder picked with the File System Access API (or the origin's private file system), `memoryStore()` into memory ([formats](docs/concepts/formats.md#a-region-written-to-a-store-browser-engine)). Its diff --git a/docs/concepts/caching.md b/docs/concepts/caching.md index 7a95108..8f20202 100644 --- a/docs/concepts/caching.md +++ b/docs/concepts/caching.md @@ -232,9 +232,9 @@ A chunk only materialize asked for is not kept once written (the kept chunks are clients that fetch a chunk again), so a large region does not push what the viewer is looking at out of the 256 MB; one a client asked for too, while it was on its way, is kept as any other. So materialize asks for each chunk once: a box written as exactly the box -(`extent: "box"`) that does not start on chunk edges has each of its chunks stitched from -the served chunks it overlaps as they come, rather than asking a served chunk again for -each of the box's chunks it overlaps. A staged view's stage outputs are kept as usual +(`extent: "box"`, the default) that does not start on chunk edges has each of its chunks +stitched from the served chunks it overlaps as they come, rather than asking a served +chunk again for each of the box's chunks it overlaps. A staged view's stage outputs are kept as usual (512 MB), so writing a large region does cycle through those. ## What has been computed diff --git a/docs/concepts/formats.md b/docs/concepts/formats.md index f28bd16..52941a0 100644 --- a/docs/concepts/formats.md +++ b/docs/concepts/formats.md @@ -90,7 +90,9 @@ part way is still the dataset, zero where no chunk was written yet; then the chu levels asked for (all by default), several at once, each of nothing but zeros left out (`--keep-empty` writes them). `--box` takes level 0's voxels along the last axes, `lo:hi` each comma separated in C order (`0,0,0:64,512,512`), and every level writes the same -region on its own voxels, rounded out to whole chunks. From Python, +region on its own voxels, rounded out to whole chunks, in an array of the level's whole +shape: what the browser engine writes with `extent: "level"`. Writing exactly the box +(the browser engine's default, `extent: "box"`) is not supported here yet. From Python, `chunkmirage.materialize(pipeline, where, format=, levels=, box=, compressor=, skip_empty=, progress=)` does the same for any pipeline, another dataset's chain read through `dataset://` among them (`registry.get(name)`), and returns the counts @@ -108,24 +110,34 @@ did and lands in the same place: the group lists the level written as its one da `0`, with that level's scale; the array has the chunk shape and dtype served, and chunks compressed with gzip unless asked not to be. What the array covers is the `extent`: -* `"level"` (the default): the level's whole shape, with the level's translation, so a box - keeps its place in it; the chunks of the box are written, rounded out to whole chunks, - each the chunk the engine serves. -* `"box"`: exactly the box, `hi − lo` voxels (clipped to the level, rounded out to whole - voxels, not to chunks), with the level's translation moved by `lo` voxels (the - translation is where voxel 0's centre is, so it lands where the view's voxel `lo` does). - Its chunks have the served chunk shape, counted from `lo`. Where `lo` is not on the - served chunks' edges, each of the box's chunks is stitched from the served chunks it - overlaps, up to two along each axis, each served chunk computed once: they are walked in - C order and each is copied into the box's chunks it overlaps, a box chunk written once - the last of them is in, so only the box's chunks begun and not yet complete are held, at - most a slab of them (one chunk thick along the first axis on which `lo` is off the - chunk edges), a row and one more; none where `lo` is on chunk edges, the box's chunks - then the served chunks as they are. The chunks at the box's far edge overhang it and are - stored whole, zero (the fill value) past the box, as zarr stores a chunk at an array's - edge (zarr v3's regular grid: border chunks have the full chunk size, the fill value - recommended outside the array; zarr v2: every chunk has the same shape). A box of the - whole level is the same store as the level's extent, byte for byte. +* `"box"` (the default): exactly the box, `hi − lo` voxels (clipped to the level, rounded + out to whole voxels, not to chunks), with the level's translation moved by `lo` voxels + (the translation is where voxel 0's centre is, so it lands where the view's voxel `lo` + does). Its chunks have the served chunk shape, counted from `lo`. The chunks at the box's + far edge overhang it and are stored whole, zero (the fill value) past the box, as zarr + stores a chunk at an array's edge (zarr v3's regular grid: border chunks have the full + chunk size, the fill value recommended outside the array; zarr v2: every chunk has the + same shape). +* `"level"`: the level's whole shape, with the level's translation, so a box keeps its + place in it, as several boxes written into one store need; the chunks of the box are + written, rounded out to whole chunks, each the chunk the engine serves. + +A box of the whole level is the same store either way, byte for byte. Where a box's `lo` is +off the served chunks' edges along an axis, the box splits the served chunks there: each of +its chunks is stitched from the served chunks it overlaps, up to two along each such axis, +and each served chunk is computed once. The served chunks are walked axis by axis, the axes +the box does not split outermost, then the split ones, the one with the most of the box's +chunks first, and each is copied into the box's chunks it overlaps, a box chunk written +once the last of them is in. Only the box's chunks begun and not yet complete are held: a +slab of them across the split axes but the outermost, a row and one more. With the box's +chunks n₁ ≥ n₂ ≥ n₃ along three split axes that is n₂ × n₃ + n₃ + 1; along two, n₂ + 1; +along one, 1; none where `lo` is on the chunk edges on every axis, the box's chunks then +the served chunks as they are. A box chunk none of whose served chunks holds a value but +zero takes no memory. So a 512 × 2048 × 2048 box of 64³ chunks off the chunk edges +everywhere (8 × 32 × 32 chunks) holds 32 × 8 + 8 + 1 = 265 chunks, which for 3-channel +float32 chunks (3 MiB) is about 800 MiB; with `lo` on a chunk edge along z, 32 + 1. No +order of the axes holds fewer; putting `lo` on the chunk edges along an axis is what +shrinks it further. A chunk of nothing but zeros is left out (zarr reads a missing chunk as the fill value, 0), unless asked to be written. A channel axis the view's ops add is the first axis and diff --git a/docs/design.md b/docs/design.md index dc1c71c..ae53a9b 100644 --- a/docs/design.md +++ b/docs/design.md @@ -526,10 +526,11 @@ as text). What such a page gets, in the order it uses it: each file written whole through `createWritable()`, and `memoryStore()` one of a `Map`, for tests and for zipping a small result. The options: the view's `level` (the finest by default), a voxel box `lo`, `hi` on it (the whole level by default; clipped to the level), - `extent` (`"level"`, the default: the array is the level's whole shape and the box's - chunks are written, rounded out to whole chunks; `"box"`: the array is exactly the box, - translated to its first voxel, its chunks the served chunk shape counted from `lo`, each - stitched from the served chunks it overlaps), `format` (`"zarr2"`, OME-Zarr 0.4, the + `extent` (`"box"`, the default: the array is exactly the box, translated to its first + voxel, its chunks the served chunk shape counted from `lo`, each stitched from the served + chunks it overlaps, since a page passing a box expects an array of that box; `"level"`: + the array is the level's whole shape and the box's chunks are written, rounded out to + whole chunks, for several boxes written into one store), `format` (`"zarr2"`, OME-Zarr 0.4, the default, since more tools read it; or `"zarr3"`, OME-Zarr 0.5), `compression` (`"gzip"`, the default, or `null`), `skipEmpty` (a chunk of nothing but zeros left out, by default), `concurrency`, `signal` (an `AbortSignal`), `onProgress` (`done`, `total`, `written`, @@ -543,9 +544,12 @@ as text). What such a page gets, in the order it uses it: does not start on the served chunks' edges is not computed again on the box's own grid: the engine keeps no chunk only materialize asked for, so asking a served chunk again for each of the box's chunks it overlaps would compute it up to eight times; instead the - served chunks are each asked for once, in C order, and copied into the box's chunks they - overlap, a box chunk written once complete, so about a slab of the box's chunks is held - at once. A single level rather than the pyramid: the engine downsamples nothing itself, + served chunks are each asked for once and copied into the box's chunks they overlap, a + box chunk written once complete, so about a slab of the box's chunks is held at once. The + walk puts the axes the box does not split outermost and the split axis with the most of + the box's chunks next, so the slab is the one across the other split axes, the smallest + any order of the axes gives (a wide, shallow box walked in C order would hold a slab + across its two wide axes). A single level rather than the pyramid: the engine downsamples nothing itself, and every level would be the model run again on its own grid. The metadata is written first, so a store whose writing stopped is still an OME-Zarr, zero where no chunk was written. The chunks come from the same producer the routes use, through the same queue, but in the diff --git a/docs/reference/cli.md b/docs/reference/cli.md index a17441d..d9a3d7c 100644 --- a/docs/reference/cli.md +++ b/docs/reference/cli.md @@ -66,7 +66,7 @@ batch's progress on stderr. | `--op`, `-o` | none | op spec, as for `serve` (repeatable) | | `--format` | `zarr3` | `zarr3` (OME-Zarr 0.5), `zarr` (zarr v2, OME-Zarr 0.4, consolidated) or `n5` | | `--level`, `-l` | every level | a level whose chunks to write (repeatable); every level's metadata is written | -| `--box` | all of it | level 0's voxels along the last axes, `lo:hi` each comma separated in C order; every level writes the same region, rounded out to whole chunks | +| `--box` | all of it | level 0's voxels along the last axes, `lo:hi` each comma separated in C order; every level writes the same region, rounded out to whole chunks, in an array of the level's whole shape (the browser engine's exact box, `extent: "box"`, is not supported here yet) | | `--chunk`, `--select` | as for `serve` | output chunk shape; non-spatial axes pinned | | `--compressor` | blosc (zarr), gzip (n5) | as the format takes it | | `--keep-empty` | off | also write chunks of nothing but zeros (left out otherwise: a reader takes a missing chunk as the fill value) | diff --git a/web/src/engine.ts b/web/src/engine.ts index f5ac1c6..df14a51 100644 --- a/web/src/engine.ts +++ b/web/src/engine.ts @@ -502,10 +502,10 @@ export class Engine { /** Write a region of `view` into `store` as an OME-Zarr of its own (materialize.ts): one * level of the view (`options.level`, the finest by default) over a box of its voxels * (`lo`, `hi`; the whole level by default), as zarr v2 with OME-Zarr 0.4 or zarr v3 with - * OME-Zarr 0.5, gzipped, its chunks of nothing but zeros left out. With `extent` "level" - * (the default) the array is the level's whole shape and the box's chunks are written, - * rounded out to whole chunks; with "box" it is exactly the box, translated to its first - * voxel, each of its chunks stitched from the view's it overlaps. The group and the array + * OME-Zarr 0.5, gzipped, its chunks of nothing but zeros left out. With `extent` "box" + * (the default) the array is exactly the box, translated to its first voxel, each of its + * chunks stitched from the view's it overlaps; with "level" it is the level's whole shape + * and the box's chunks are written, rounded out to whole chunks. The group and the array * are written first, so the store reads as the view from the start, zero where no chunk is * yet; then the chunks, each of the view's computed once, as a client's request for it is * (kept, or shared with one on its way, and kept after) but in the background: behind diff --git a/web/src/layout.ts b/web/src/layout.ts index 97d4af7..d93d0fa 100644 --- a/web/src/layout.ts +++ b/web/src/layout.ts @@ -12,10 +12,6 @@ const OME_UNIT: Record = { nm: "nanometer", um: "micrometer", m: export const V2_DTYPE: Record = { uint8: "|u1", int8: "|i1", uint16: " = { - uint8: 1, int8: 1, uint16: 2, int16: 2, uint32: 4, int32: 4, uint64: 8, int64: 8, float32: 4, float64: 8, -}; export const isTime = (a: ViewAxis) => a.unit in TO_SECONDS; /** A level of a view as one of a group's datasets: its path in the group, and which of the diff --git a/web/src/materialize.ts b/web/src/materialize.ts index 74d1ed9..5916a2f 100644 --- a/web/src/materialize.ts +++ b/web/src/materialize.ts @@ -8,7 +8,8 @@ // written, gzip, the array's chunks stitched from the view's for a box that does not start // on chunk edges, and the loop that has the chunks made and writes them, from any producer // (the engine's queue on the page, a fake one in the tests). -import { array2, array3, attrs2, chunkKey, dimensionNames, group3, ITEM_BYTES } from "./layout.ts"; // with its extension: Node's runner resolves no other +import { TYPED } from "./asread.ts"; // with its extension: Node's runner resolves no other +import { array2, array3, attrs2, chunkKey, dimensionNames, group3 } from "./layout.ts"; import type { ViewInfo } from "./types.ts"; /** Where materialize writes: files by "/"-separated path relative to the store's root. */ @@ -92,14 +93,16 @@ export interface MaterializeOptions { * whole level. Clipped to the level, and rounded out to whole voxels (to whole chunks, with * the level's `extent`). */ lo?: number[]; hi?: number[]; - /** What the written array covers. "level" (the default): the level's whole shape, with the - * level's translation, so the box keeps its place in it; the chunks the box touches are - * written, rounded out to whole chunks. "box": exactly the box, not rounded to chunks: its - * shape hi − lo, its translation the level's moved by lo voxels, cut in the view's chunk - * shape counted from lo. Where the box does not start on the view's chunk edges, each of - * its chunks is stitched from the view's chunks it overlaps (up to 2 along each axis), each - * of those computed once; the chunks at its far edge overhang it and are stored whole, the - * fill value past the box, as zarr stores a chunk at an array's edge. */ + /** What the written array covers. "box" (the default): exactly the box, not rounded to + * chunks: its shape hi − lo, its translation the level's moved by lo voxels, cut in the + * view's chunk shape counted from lo. Where the box does not start on the view's chunk + * edges, each of its chunks is stitched from the view's chunks it overlaps (up to 2 along + * each axis), each of those computed once; the chunks at its far edge overhang it and are + * stored whole, the fill value past the box, as zarr stores a chunk at an array's edge. + * "level": the level's whole shape, with the level's translation, so the box keeps its + * place in it (several boxes written into one store, say); the chunks the box touches are + * written, rounded out to whole chunks. A box of the whole level is the same store, byte + * for byte, either way. */ extent?: "level" | "box"; /** "zarr2" (OME-Zarr 0.4: .zgroup, .zattrs, .zarray, "/" as the dimension separator) or * "zarr3" (OME-Zarr 0.5, zarr.json); default "zarr2", which more tools read. */ @@ -162,15 +165,16 @@ export function chunkRange(shape: number[], chunk: number[], lo?: number[], hi?: return { lo: box.lo.map((x, a) => Math.floor(x / chunk[a])), hi: box.hi.map((x, a) => Math.ceil(x / chunk[a])) }; } -/** The chunk indices [lo, hi), in C order (the last axis fastest). */ -export function* chunkIndices(lo: number[], hi: number[]): Generator { +/** The chunk indices [lo, hi), in C order (the last axis fastest), or in `order`, the axes + * from the outermost to the innermost (the last of them fastest). */ +export function* chunkIndices(lo: number[], hi: number[], order: number[] = lo.map((_, a) => a)): Generator { if (hi.some((h, a) => h <= lo[a])) return; const at = [...lo]; for (;;) { yield [...at]; - let a = at.length - 1; - while (a >= 0 && ++at[a] >= hi[a]) { at[a] = lo[a]; a--; } - if (a < 0) return; + let k = order.length - 1; + while (k >= 0 && ++at[order[k]] >= hi[order[k]]) { at[order[k]] = lo[order[k]]; k--; } + if (k < 0) return; } } @@ -222,14 +226,14 @@ const count = (lo: number[], hi: number[]) => prod(lo.map((l, a) => hi[a] - l)); * whole shape (so a box keeps its place in it) or the box's, the chunks and dtype served, * compressed as asked; and the chunks of the box. */ export function region(name: string, v: ViewInfo, chunk: number[], options: MaterializeOptions = {}): Region { - const { level = 0, format = "zarr2", compression = "gzip", extent = "level", attributes = {} } = options; + const { level = 0, format = "zarr2", compression = "gzip", extent = "box", attributes = {} } = options; if (!Number.isInteger(level) || !v.levels[level]) throw new Error(`${name} has no level ${level}: it has ${v.levels.length}`); if (format !== "zarr2" && format !== "zarr3") throw new Error(`format ${JSON.stringify(format)}: "zarr2" or "zarr3"`); if (compression !== "gzip" && compression !== null) throw new Error(`compression ${JSON.stringify(compression)}: "gzip" or null`); if (extent !== "level" && extent !== "box") throw new Error(`extent ${JSON.stringify(extent)}: "level" or "box"`); const levelShape = v.levels[level].shape, box = voxelBox(levelShape, options.lo, options.hi); const { lo, hi } = chunkRange(levelShape, chunk, box.lo, box.hi); // the view's chunks the box needs - const whole = extent === "level", item = ITEM_BYTES[v.out]; + const whole = extent === "level", item = TYPED[v.out]?.BYTES_PER_ELEMENT; if (!whole && !item) throw new Error(`${name} is ${v.out}, which a box is not cut from`); const at = whole ? levelShape.map(() => 0) : box.lo, shape = whole ? levelShape : box.hi.map((h, a) => h - box.lo[a]); const first = whole ? lo : shape.map(() => 0), last = whole ? hi : shape.map((n, a) => Math.ceil(n / chunk[a])); @@ -255,6 +259,28 @@ export function region(name: string, v: ViewInfo, chunk: number[], options: Mate }; } +/** The axes along which region `r`'s array chunks are stitched from two of the view's + * chunks: those along which the array's first chunk spans a chunk edge of the view's (the + * box starts off the view's chunk edges there, and does not lie within one of its chunks). */ +export function splitAxes(r: Region): number[] { + return r.chunk.flatMap((c, a) => { + const from = r.at[a] + r.first[a] * c, to = r.at[a] + Math.min((r.first[a] + 1) * c, r.shape[a]); + return Math.floor(from / c) !== Math.floor((to - 1) / c) ? [a] : []; + }); +} + +/** The order region `r`'s served chunks are walked in, its axes from the outermost to the + * innermost: first the axes it does not split (`splitAxes`), in C order, then those it + * splits, the one with the most of the array's chunks first (C order among equals). Of the + * orders of the axes, it holds the fewest array chunks at once (`stitcher`): a slab across + * the split axes but the outermost, so the longest split axis is the one walked slab by + * slab. The level's extent splits none: C order, as the view's chunks are numbered. */ +export function walkOrder(r: Region): number[] { + const split = splitAxes(r), count = (a: number) => r.last[a] - r.first[a]; + const rest = r.chunk.map((_, a) => a).filter((a) => !split.includes(a)); + return [...rest, ...split.sort((a, b) => count(b) - count(a) || a - b)]; +} + /** The bytes of a chunk of region `r`'s array (or of the view's: they are cut alike). */ export const chunkBytes = (r: Region) => r.channels * prod(r.chunk) * r.item; @@ -271,10 +297,11 @@ export type Completed = [index: number[], bytes: Uint8Array | null]; * chunks it overlaps that holds a value not zero comes (none, if they are all zero), each * one's part copied in, the rest the fill value; it is complete, and forgotten here, once all * of them have come. The array's chunks begun and not complete are `open`: with the view's - * chunks added in C order, at most a slab of them (one chunk thick along the first axis on - * which the box does not start on a chunk edge, the box's chunks across the axes after it), - * a row and one more, so a row and one where the box starts on chunk edges along z alone, - * one along z and y, and none along every axis; a few more where they come out of order. */ + * chunks added in `walkOrder`, at most a slab of them across the split axes but the + * outermost (`splitAxes`), a row along the innermost and one more. With the box's chunks + * n1 ≥ n2 ≥ n3 along three split axes, n2 × n3 + n3 + 1; along two, n1 ≥ n2, n2 + 1; along + * one, 1; none if it splits none (the box starts on the view's chunk edges). A few more + * where they come out of order. */ export function stitcher(r: Region): { add: (index: number[], bytes: Uint8Array) => Completed[]; readonly open: number } { const C = r.chunk, n = C.length, size = chunkBytes(r); const open = new Map(); @@ -345,8 +372,9 @@ export type Produce = (index: number[], signal: AbortSignal) => Promise String((e as Error)?.message ?? e); /** Write region `r` into `store`: its metadata, then its chunks. The view's chunks the box - * covers are asked of `produce` once each, in C order (the last axis fastest), - * `concurrency` at once (one by default); each, as it comes, completes the array's chunks it + * covers are asked of `produce` once each, in `walkOrder` (C order where the box splits no + * chunk of the view's, as with the level's extent), `concurrency` at once (one by default); + * each, as it comes, completes the array's chunks it * is the last of (`stitcher`: with the level's extent, or a box on the view's chunk edges, * just itself), and those are written, gzipped if `r` says so, those of nothing but zeros * skipped (`skipEmpty`, the default). Only the array's chunks begun and not yet complete are @@ -381,7 +409,7 @@ export async function writeChunks(store: WritableStore, r: Region, produce: Prod report(); // a chunk of the view named as the array's chunk it is, with the level's extent; as the engine names it, with the box's const named = (index: number[]) => (r.extent === "level" ? `chunk ${r.key(index)} of ${r.name}` : `chunk ${r.name}/${r.level}/${index.join(".")}`); - const stitched = stitcher(r), indices = chunkIndices(r.lo, r.hi); + const stitched = stitcher(r), indices = chunkIndices(r.lo, r.hi, walkOrder(r)); const lane = async () => { for (let next = indices.next(); !next.done && !stop.signal.aborted; next = indices.next()) { const index = next.value; diff --git a/web/test/materialize.test.ts b/web/test/materialize.test.ts index bde35c7..003ebc2 100644 --- a/web/test/materialize.test.ts +++ b/web/test/materialize.test.ts @@ -4,13 +4,17 @@ // producer, with empty chunks skipped, progress, abort and failure; the stores; and the // result read back by zarrita, a zarr implementation of its own, as a reader would. Then the // same for a box written as exactly the box (extent "box"): its metadata, its chunks -// stitched from the view's (each of those made once), its edge chunks, and the slab of -// chunks held at most. Run by `npm test`. +// stitched from the view's (each of those made once), its edge chunks, the order the view's +// chunks are walked in and the slab of chunks held at most. Run by `npm test`. The level's +// extent is asked for by name where a test is of it: a box is the default. import { test } from "node:test"; import assert from "node:assert/strict"; import * as zarr from "zarrita"; import { array3, chunkKey, group3 } from "../src/layout.ts"; -import { allZero, chunkIndices, chunkRange, directoryStore, gzip, memoryStore, region, stitcher, voxelBox, writeChunks, type MaterializeOptions, type MaterializeProgress, type Produce, type WritableStore } from "../src/materialize.ts"; +import { + allZero, chunkIndices, chunkRange, directoryStore, gzip, memoryStore, region, splitAxes, stitcher, voxelBox, walkOrder, writeChunks, + type MaterializeOptions, type MaterializeProgress, type Produce, type Region, type WritableStore, +} from "../src/materialize.ts"; import type { ViewInfo } from "../src/types.ts"; const nm = (name: string) => ({ name, unit: "nm" }); @@ -67,6 +71,9 @@ test("the chunks of a range in C order, the last axis fastest", () => { assert.deepEqual([...chunkIndices([0, 1], [2, 3])], [[0, 1], [0, 2], [1, 1], [1, 2]]); assert.deepEqual([...chunkIndices([1, 1, 1], [2, 2, 2])], [[1, 1, 1]]); assert.deepEqual([...chunkIndices([0, 0], [0, 3])], []); + // in another order of the axes, the last of it fastest + assert.deepEqual([...chunkIndices([0, 1], [2, 3], [1, 0])], [[0, 1], [1, 1], [0, 2], [1, 2]]); + assert.deepEqual([...chunkIndices([0, 0, 0], [2, 1, 2], [2, 0, 1])].map((i) => i.join("")), ["000", "100", "001", "101"]); }); test("zarr v2: .zgroup, .zattrs with OME 0.4 multiscales listing the one level at 0, its .zarray gzipped and .zattrs", () => { @@ -93,7 +100,7 @@ test("zarr v2: .zgroup, .zattrs with OME 0.4 multiscales listing the one level a }); test("zarr v3: zarr.json with OME 0.5 multiscales, the level's zarr.json with the gzip codec after bytes", () => { - const r = region("mito", view, chunk, { format: "zarr3", lo: [2, 0, 4], hi: [5, 3, 7], attributes: { provenance: 1 } }); + const r = region("mito", view, chunk, { format: "zarr3", extent: "level", lo: [2, 0, 4], hi: [5, 3, 7], attributes: { provenance: 1 } }); assert.deepEqual(r.files.map(([p]) => p), ["zarr.json", "0/zarr.json"]); const [group, array] = r.files.map(([, b]) => json(b)); assert.deepEqual(group, { @@ -162,10 +169,10 @@ test("gzip round trips through CompressionStream and DecompressionStream; zeros }); for (const format of ["zarr2", "zarr3"] as const) { - test(`${format}: the box's chunks written gzipped, the empty one skipped, and read back by zarrita as the view`, async () => { + test(`${format}, the level's extent: the box's chunks written gzipped, the empty one skipped, and read back by zarrita as the view`, async () => { const store = memoryStore(), seen: MaterializeProgress[] = []; // z voxels 1 to 3 of the whole y, x: chunks z 0 and 1, the second of them empty - const got = await writeChunks(store, region("v", view, chunk, { format, lo: [1, 0, 0], hi: [3, 6, 7] }), producer(["1.0.1"]), + const got = await writeChunks(store, region("v", view, chunk, { format, extent: "level", lo: [1, 0, 0], hi: [3, 6, 7] }), producer(["1.0.1"]), { concurrency: 3, onProgress: (p) => seen.push(p) }); const keys = files(store).filter((f) => !f.endsWith("json") && !f.includes(".z")); const at = (i: string) => (format === "zarr3" ? `0/c/${i}` : `0/${i}`); @@ -196,7 +203,7 @@ for (const format of ["zarr2", "zarr3"] as const) { test("empty chunks are written too when asked, and uncompressed ones as they are", async () => { const store = memoryStore(); - const got = await writeChunks(store, region("v", view, chunk, { compression: null, lo: [2, 0, 4], hi: [4, 4, 8] }), producer(["1.0.1"]), { skipEmpty: false }); + const got = await writeChunks(store, region("v", view, chunk, { extent: "level", compression: null, lo: [2, 0, 4], hi: [4, 4, 8] }), producer(["1.0.1"]), { skipEmpty: false }); assert.deepEqual(got, { chunks: 1, written: 1, skipped: 0, bytes: 64, aborted: false }); assert.deepEqual(store.files.get("0/1/0/1"), new Uint8Array(32 * 2)); }); @@ -251,7 +258,7 @@ test("a chunk that fails rejects naming it and why, once the writes under way ha if (index.join(".") === "0.0.1") { await tick(); throw new Error("the model ran out of GPU memory"); } return fakeChunk(index); }; - await assert.rejects(writeChunks(store, region("v", view, chunk), produce, { concurrency: 2 }), (e: Error) => { + await assert.rejects(writeChunks(store, region("v", view, chunk, { extent: "level" }), produce, { concurrency: 2 }), (e: Error) => { assert.equal(e.message, "chunk 0/0/0/1 of v: the model ran out of GPU memory"); assert.ok(slowWritten, "the chunk being written when the other failed was written first"); return true; @@ -395,8 +402,9 @@ test("a box's array is exactly the box: its shape hi − lo, the view's chunks c assert.deepEqual([gArray.shape, gArray.chunks], [[2, 3, 3, 5], [2, 2, 4, 4]]); assert.deepEqual(gAttrs.multiscales[0].datasets[0].coordinateTransformations[1].translation, [0, 8, 24, 16]); assert.equal(grad.key([1, 0, 1]), "0/0/1/0/1"); - // the level's own extent is as before: the level's whole shape and translation - const level = region("mito", view, chunk, { level: 1, lo: [1, 0, 2], hi: [3, 2, 4] }); + // a box is the default; the level's own extent is as before: the level's whole shape and translation + assert.deepEqual(region("v", view, chunk, { lo: [1, 3, 2], hi: [4, 6, 7] }).shape, [3, 3, 5]); + const level = region("mito", view, chunk, { level: 1, extent: "level", lo: [1, 0, 2], hi: [3, 2, 4] }); assert.deepEqual(json(level.files[2][1]).shape, [3, 3, 4]); assert.deepEqual(json(level.files[1][1]).multiscales[0].datasets[0].coordinateTransformations[1].translation, [4, 4, 4]); }); @@ -483,7 +491,7 @@ for (const format of ["zarr2", "zarr3"] as const) { test(`${format}: a box of the whole level is the level's extent, file for file and byte for byte`, async () => { for (const box of [{}, { lo: [0, 0, 0], hi: [5, 6, 7] }, { lo: [-1, -1, -1], hi: [9, 9, 9] }] as MaterializeOptions[]) { const level = memoryStore(), whole = memoryStore(); - await writeChunks(level, region("v", view, chunk, { format, ...box }), producer(["1.0.1"])); + await writeChunks(level, region("v", view, chunk, { format, extent: "level", ...box }), producer(["1.0.1"])); await writeChunks(whole, region("v", view, chunk, { format, extent: "box", ...box }), producer(["1.0.1"])); assert.deepEqual(whole.files, level.files); } @@ -506,9 +514,20 @@ test("skipEmpty: a box's chunk is left out when the view's chunks it is made of assertBox((await readArray(all, "zarr2")).data, [4, 8, 8], [1, 1, 1], zero); }); +/** The most array chunks `r`'s stitcher holds with the view's chunks added in `order` + * (walkOrder's by default), each of them made by `make`; none held at the end. */ +function held(r: Region, make: (index: number[]) => ArrayBuffer, order = walkOrder(r)): number { + const s = stitcher(r); + let most = 0; + for (const i of chunkIndices(r.lo, r.hi, order)) { s.add(i, new Uint8Array(make(i))); most = Math.max(most, s.open); } + assert.equal(s.open, 0); + return most; +} + test("a box's chunk is written once the last of the view's chunks it is made of comes; a slab of them held at most", async () => { // from (1, 1, 1) to the level's end: off the chunk edges on every axis, 8 × 9 × 10 in 4 × 3 × 3 chunks - const r = region("v", big, chunk, { extent: "box", compression: null, lo: [1, 1, 1] }), events: string[] = []; + const r = region("v", big, chunk, { compression: null, lo: [1, 1, 1] }), events: string[] = []; + assert.deepEqual([splitAxes(r), walkOrder(r)], [[0, 1, 2], [0, 1, 2]]); // z has the most chunks: C order const fake = served(big), produce: Produce = async (i, signal) => { events.push(`make ${i.join(".")}`); return fake.produce(i, signal); }; const store: WritableStore = { async put(path) { if (/^0\/\d/.test(path)) events.push(`write ${path}`); } }; await writeChunks(store, r, produce); @@ -521,24 +540,42 @@ test("a box's chunk is written once the last of the view's chunks it is made of const last = o.map((x, a) => Math.floor((r.at[a] + Math.min((x + 1) * chunk[a], r.shape[a]) - 1) / chunk[a])); assert.equal(events.slice(0, i).filter((x) => x.startsWith("make")).at(-1), `make ${last.join(".")}`, `${e} as soon as it is complete`); } - // held: off chunk edges along z, a slab of the box's chunks (3 × 3), a row (3) and one more at most - const slab = (lo: number[]) => { - const s = stitcher(region("v", big, chunk, { extent: "box", compression: null, lo })), r2 = region("v", big, chunk, { extent: "box", lo }); - let most = 0; - for (const i of chunkIndices(r2.lo, r2.hi)) { s.add(i, new Uint8Array(fake.make(i))); most = Math.max(most, s.open); } - assert.equal(s.open, 0); // all complete at the end - return most; - }; - assert.equal(slab([1, 1, 1]), 3 * 3 + 3 + 1); - assert.equal(slab([2, 1, 1]), 3 + 1); // on edges along z: a row and one - assert.equal(slab([2, 4, 1]), 1); // along z and y: one - assert.equal(slab([2, 4, 4]), 0); + // held: split along all three (4 ≥ 3 ≥ 3 chunks), a slab across the inner two (3 × 3), a row (3) and one more + const box = (lo: number[]) => region("v", big, chunk, { lo }); + assert.equal(held(box([1, 1, 1]), fake.make), 3 * 3 + 3 + 1); + assert.equal(held(box([2, 1, 1]), fake.make), 3 + 1); // on edges along z: split along y and x, a row and one + assert.equal(held(box([2, 4, 1]), fake.make), 1); // along z and y: one + assert.equal(held(box([2, 4, 4]), fake.make), 0); // along all: none, each of the view's passing through + assert.deepEqual(splitAxes(box([2, 4, 4])), []); // the view's chunks all empty: nothing held in memory but the counts (no buffer made) - const r3 = region("v", big, chunk, { extent: "box", lo: [1, 1, 1] }), s3 = stitcher(r3); + const r3 = box([1, 1, 1]), s3 = stitcher(r3); const done = [...chunkIndices(r3.lo, r3.hi)].flatMap((i) => s3.add(i, new Uint8Array(64))); assert.deepEqual([done.length, done.every(([, b]) => b === null)], [r3.total, true]); }); +test("the view's chunks are walked with the split axis of the most chunks outermost, so a wide, shallow box holds a thin slab", async () => { + // a level 4 × 40 × 40 and a box from (1, 1, 1): 3 × 39 × 39 voxels, 2 × 10 × 10 chunks, split along every axis + const wide: ViewInfo = { ...view, levels: [{ shape: [4, 40, 40], voxel: [8, 8, 8], origin: [0, 0, 0] }] }; + const fake = served(wide), r = region("v", wide, chunk, { compression: null, lo: [1, 1, 1] }); + assert.deepEqual([r.last, splitAxes(r), walkOrder(r)], [[2, 10, 10], [0, 1, 2], [1, 2, 0]]); // y, then x, z fastest + // held: a slab across x and z (10 × 2), a row along z (2) and one more; in C order, a slab across y and x + assert.equal(held(r, fake.make), 10 * 2 + 2 + 1); + assert.equal(held(r, fake.make, [0, 1, 2]), 10 * 10 + 10 + 1); + assert.equal(held(r, fake.make, [0, 2, 1]), 10 * 10 + 10 + 1); // the fewest chunks outermost: the most held + // the axes it does not split go outermost: on chunk edges along y, then x (10 chunks), then z (2) + const edge = region("v", wide, chunk, { lo: [1, 4, 1] }); + assert.deepEqual([splitAxes(edge), walkOrder(edge)], [[0, 2], [1, 2, 0]]); + assert.equal(held(edge, fake.make), 2 + 1); + assert.ok(held(edge, fake.make, [0, 1, 2]) > 9 * 10); // C order: a slab across y and x + // writeChunks asks in that order, each once, and the box reads back as the view + const store = memoryStore(); + await writeChunks(store, region("v", wide, chunk, { lo: [1, 1, 1] }), fake.produce, { concurrency: 2 }); + assert.deepEqual(fake.asked.slice(0, 4), ["0.0.0", "1.0.0", "0.0.1", "1.0.1"]); + assert.deepEqual(fake.asked, [...chunkIndices(r.lo, r.hi, [1, 2, 0])].map((i) => i.join("."))); + assert.ok([...fake.calls.values()].every((n) => n === 1)); + assertBox((await readArray(store, "zarr2")).data, [3, 39, 39], [1, 1, 1]); +}); + test("a box's abort: no chunk asked for or written after it, those asked for given up; the chunks written stay", { timeout: 5000 }, async () => { const store = memoryStore(), stop = new AbortController(), asked: string[] = [], given: AbortSignal[] = [], fake = served(big); let atAbort = -1;