Browser engine: Engine.materialize writes exactly the box, stitched from the served chunks (extent: "box", the default)

[davidackerman]

Why

cellmap-flow's browser app (its Run tab) uses Engine.materialize to save a model's prediction over a box the user draws. Until now the saved array always had the whole level's shape. Only the box's chunks were present, the box was rounded out to whole chunks, and the level's translation was kept. A 300³ box on a large level became a level-sized array with a few chunks in it, plus up to a chunk of extra voxels on each side. The user asked: "write out only the box, also don't do it in whole chunks, you can just truncate the writing of the last chunks."

There are two commits:

• 431499b adds the option.
• 2217e2b applies the review of it: "box" becomes the default, the walk order is chosen to keep the buffered slab small, and the byte widths come from TYPED.

The API

One new option, in web/src/materialize.ts:

interface MaterializeOptions {
  // ...level, lo, hi, format, compression, skipEmpty, concurrency, signal, onProgress, attributes, as before
  /** "box" (the default): exactly the box. "level": the level's whole shape, the box's chunks in it. */
  extent?: "box" | "level";
}

• "box" (the default). A page that passes a box expects an array of that box. The array is hi − lo voxels on the spatial axes. The box is clipped to the level and rounded out to whole voxels, not to chunks. Its translation is the level's, moved by lo voxels. Its chunks have the served chunk shape, counted from lo. Channel axes added by the ops stay whole and first.
• "level". The array has the level's whole shape and the level's translation, and the box's chunks are written into it, rounded out to whole chunks. This is what every call did before, and it is what several boxes written into one store need.
• A box of the whole level is the same store either way, byte for byte. That includes passing no lo/hi (tested).

MaterializeProgress gains computed and needed: the served chunks computed so far, out of those the box covers.

• done, total, written and bytes still count the written array's chunks.
• onProgress is called once after the metadata, then once per served chunk, after the array chunks it completes have been written. With "level" that is once per chunk, as before.
• Every count only grows. With "box", done can trail computed by about a slab, so computed/needed lets a progress bar move during the first slab.
• Materialized is unchanged.

A served chunk that fails during a "box" write is named the way the engine's own log names it: chunk v/0/2.0.0: <why>. "level" keeps chunk 0/0/0/1 of v: <why>.

Smaller additions:

• materialize.ts exports voxelBox, splitAxes, walkOrder, stitcher, chunkBytes and Completed.
  • chunkIndices(lo, hi, order?) takes an axis order.
  • chunkRange now uses voxelBox and gives the same results.
  • Region gains the array's geometry.
• layout.ts gains Dataset.lo?, translation(v, level, lo?), array3(..., { shape }) and array2(..., compressor, shape). They default to the level's shape, so the routes are unchanged.
• The width of a value comes from asread.ts's TYPED (BYTES_PER_ELEMENT).

What it writes

• Translation. It keeps its meaning: the centre of voxel 0, origin + lo × voxel (formats.md says "a translation is where voxel 0's centre is"; ome.ts does the same for a region). The scale is the level's, and the dataset path stays "0".
• Edge chunks. The chunks at the box's far edge overhang it. They are stored at the full chunk shape, padded with 0 (the fill value) past the box.
  • zarr v3's regular grid says: "Chunks at the border of an array always have the full chunk size", and "it is recommended to use the current fill value for elements outside the bounds".
  • zarr v2 says all chunks in an array have the same shape, and what lies outside the array is undefined.
• Reader checks. Box stores in zarr2 and zarr3, gzip and raw, off the chunk edges on every axis, were read back equal by zarrita (in the tests), and by tensorstore and zarr-python 3.4.1 outside them. The translation was [12, 28, 20] for origin 4, voxel 8, lo = [1, 3, 2].
• Truncated chunks are refused. With one raw edge chunk cut down to its in-box voxels, zarr-python ("cannot reshape array of size 2 into shape (2,4,4)") and tensorstore ("Not enough data") both refuse it. So "truncating the last chunks" happens in the array's shape, not in the chunk bytes.
• Past the level's edge, not forced to zero. When a box ends at the level's edge, a chunk passed through unchanged keeps the served bytes past that edge, as the level's extent writes them. Served chunks are already zero there.

Stitching: each served chunk computed once

A chunk that only materialize asked for is not kept in the engine's chunk cache (ac41da8). Asking for the served chunks again for each box chunk they overlap would therefore compute each one up to 8 times. Instead:

1. One walk. writeChunks walks the served chunks the box covers once each, in walkOrder (below), through the same produce. On the page that is the existing background claim: background claims, the bounded queue, slots−1 and promotion by a client are untouched. Background jobs at one level run in submission order, so the queue follows the walk.
2. Scatter. stitcher(r).add(index, bytes) copies the served chunk into each box chunk it overlaps, up to 2 per axis, a row at a time. A box chunk counts the served chunks it is made of. When the last one arrives, the box chunk is returned and written at once (gzip, put), then forgotten.
3. Pass-through. A box chunk that is exactly one served chunk (the same voxels, and ending where the level ends if it overhangs) is written as the served bytes, without copying. That covers every chunk with "level", and every chunk of a box that starts on chunk edges.
4. Lazy buffers. A box chunk's buffer is allocated only when one of its served chunks holds a value other than zero. A box chunk made only of empty chunks takes no memory, and is skipped with skipEmpty.

Walk order and the bound on memory

Only box chunks that have been started but are not complete are held, plus the served chunks in flight.

The walk order. The served chunks are walked axis by axis:

• First, the axes the box does not split, in C order. "Split" means lo is off the chunk edges there; a box lying within one chunk along an axis doesn't split it.
• Then the split axes, the one with the most of the box's chunks first, keeping C order among equal counts.

The bound. The most box chunks held is a slab across the split axes except the outermost one, plus a row and one more. Let n₁ ≥ n₂ ≥ n₃ be the box's chunk counts along the split axes. Then:

Split axes	Most box chunks held
Three	n₂ × n₃ + n₃ + 1
Two	n₂ + 1
One	1
None	0 (every chunk passes through)

concurrency > 1 can add a few more, because chunks arrive out of order. The level's extent splits nothing, so it is walked in C order as before. The tests assert the count held, exactly, for each alignment and for several orders.

A note on the review's wording. The review asked for "the thinnest misaligned axis outermost" to keep the slab small. The slab spans the axes other than the outermost split axis, so the smallest slab comes from putting the split axis with the most chunks outermost. Putting the one with the fewest outermost is the worst order. I implemented the stated goal, the smallest slab. For 3-channel float32 affinities in 64³ chunks (3 MiB each), with the box off the chunk edges on every axis:

Box (z × y × x)	Chunks	C order	Fewest outermost	This PR	lo on a z chunk edge
512 × 2048 × 2048	8 × 32 × 32	1057 = 3.1 GiB	1057 = 3.1 GiB	265 = 795 MiB	33 ≈ 99 MiB
2048 × 2048 × 512	32 × 32 × 8	265 = 795 MiB	1057 = 3.1 GiB	265 = 795 MiB	9 ≈ 27 MiB

So the reorder fixes a wide, shallow box. A box split on all three axes still holds about a cross-section of the box's chunks, about 800 MiB in the review's 2048 × 2048 × 512 example, whatever the order. No order of the axes holds fewer. What shrinks it further is putting lo on chunk edges along an axis. Box chunks that are all background take nothing, so these figures are upper bounds.

What changes for existing callers

• The default is now "box". A call that passes lo/hi without extent now gets an array of exactly the box. A call with neither gets the same store as before, byte for byte. Pass extent: "level" for the old array. cellmap-flow, the only known caller, wants "box".
• "level" is today's write. Its metadata, chunk keys and bytes, skip rule, abort and failure behaviour, error messages and C-order walk are all unchanged.
• Engine.materialize's code is unchanged, apart from its doc comment. Producer, queue, the kept-chunks rule and background scheduling are as before.
• Progress objects gain two fields. The existing fields have the same values at the same moments.
• The Python side still writes the level's extent only. chunkmirage materialize --box and chunkmirage.materialize(box=) round out to whole chunks in an array of the level's shape. formats.md and cli.md now say the exact box is not supported there yet. The same extent option, defaulting to "box", is planned there after this merges.

Docs

• docs/concepts/formats.md:
  • The browser-engine materialize section covers both extents, the box's shape and translation, edge chunks with the specs' wording, byte-for-byte equality for a whole-level box, the walk order and bound with the example, and what progress counts.
  • The Python section says it writes the level's extent only, for now.
• docs/design.md: the options list, why each served chunk is asked for once, and the walk order.
• docs/concepts/caching.md: next to the not-kept rule, why materialize asks for each chunk once.
• docs/reference/cli.md: the --box row says the exact box is not supported there yet.
• CHANGELOG.md: the unreleased Engine.materialize bullet is amended.
• Doc comments in materialize.ts, layout.ts and engine.ts.

Tests

web/test/materialize.test.ts has 14 new tests, so npm test goes from 56 to 70. Existing tests of the old behaviour now pass extent: "level" by name. Their progress checks include the two new fields, and the C-order test also covers a different axis order.

• Metadata. Box metadata in both formats: shape, chunk grid, and translation ([20, 4, 36] on level 1, with 16 nm voxels and a 4 nm origin). Also clipping, rounding to whole voxels, the default, channel axes, and the level's extent unchanged.
• Unaligned box, both formats. It crosses chunk edges on every axis: 24 box chunks from 36 served, concurrency 3, chunks arriving out of turn.
  • Each served chunk is asked for exactly once (the producer's calls are counted).
  • zarrita reads back the served values over the box.
  • The far-corner chunk is stored at full size with zeros past the box, checked voxel by voxel.
  • Progress is monotonic: 1 + 36 reports.
• Aligned box. Stored chunks equal the served chunks byte for byte, and nothing is held. A box aligned at lo but ending inside chunks short of the level's end is cropped.
• Smaller than a chunk. A box inside one served chunk, and a box across the corner of 8.
• Whole level. The box equals the level's extent file for file and byte for byte, in both formats: with no box, with an exact one, and with one past the edges.
• skipEmpty. A box chunk made only of empty served chunks is skipped, and is written as zeros when skipping is off. A box chunk with some empty parts is written, zero there.
• Writes and walk order. Each box chunk is written right after its last served chunk arrives. The count held is exact for each alignment. A wide, shallow box (2 × 10 × 10 chunks) holds 10 × 2 + 2 + 1 = 23 in this order, against 111 in C order, and the producer's asks follow that order.
• Abort. Nothing is asked for or written after it, the served chunks asked for are given up, and what was written stays.
• Failure. It is named as the engine names it, after the write under way has ended. No other box chunk the same served chunk completed is written. A served chunk of the wrong size is refused.
• Channel axes, both formats. Two float32 channels, whole and first, read back equal.

Open questions

• About a cross-section of the box's chunks in memory for a box split on all three axes (about 800 MiB in the example above). Is that acceptable as the documented bound, with pages putting lo on chunk edges along one axis when the box is large? Or should materialize also bound it, for example by compressing partial chunks while they wait, or refusing above a budget?
• Engine.materialize itself still has no automated test, because Node can't import engine.ts (its JSON import). Its code path is unchanged here, and the parts this PR changes run in npm test over a fake producer.

Settled in the review: fractional boxes are rounded out to whole voxels; Materialized gets no computed; nothing past the level's edge is forced to zero.

Checks:

• npm run check in web/ passes: both typechecks, 70/70 tests, and the build.
• mkdocs build --strict passes.
• pytest -q tests/test_docs.py passes (6).
• ruff check src tests examples web passes.
• After the first commit, the full pytest -q passed (313 passed, 9 skipped). No Python changed after that.