From 297bf8f821ce98bd3a68db74971ba2592d302650 Mon Sep 17 00:00:00 2001 From: David Ackerman Date: Fri, 16 Oct 2026 15:51:59 -0400 Subject: [PATCH] chunkmirage materialize writes exactly the box by default, as #41's browser engine does: --extent box (chunkmirage.materialize(extent="box")) writes each level as an array of the box's shape on it, its translation moved by the box's corner so it lands where the box is, in the dataset's chunk shape counted from that corner, each chunk read from the dataset's chunks it overlaps; when the last stage keeps no chunks (no cache: true) those pass through a cache of up to 1 GiB, walked with the axes the corner is on chunk edges along first and then the one with the most chunks, so each is computed once while a slab fits (22 computed for the test's box, 84 without the cache); --extent level keeps #40's arrays of the levels' whole shapes, for several boxes in one store; tests write an off-grid box as zarr v3 and v2 and read it back in place, counting each chunk computed once; docs (formats, CLI, API) and the CHANGELOG say so, where they said the Python side wrote the level's extent only --- CHANGELOG.md | 9 ++-- docs/concepts/formats.md | 16 +++++-- docs/reference/api.md | 2 +- docs/reference/cli.md | 3 +- src/chunkmirage/cli.py | 6 ++- src/chunkmirage/materialize.py | 86 ++++++++++++++++++++++++++++++---- tests/test_materialize.py | 43 ++++++++++++++++- 7 files changed, 146 insertions(+), 19 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 4e11226..ec930f8 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -58,7 +58,10 @@ The first release, an alpha: the plugin API can still change before 0.1.0. - `chunkmirage materialize` and `chunkmirage.materialize`: a dataset, or a box of it, computed and written into a store (a directory, `gs://`, `s3://`) as a frontend serves it, zarr v3, zarr v2 or N5, so it reads back as a stored dataset of its own; the metadata - first, chunks several at once, empty ones left out. + first, chunks several at once, empty ones left out. A box is exactly the box by default + (`--extent box`: arrays of its shape, placed where it is, each of the dataset's chunks it + reaches computed once), or its chunks in arrays of the levels' whole shapes + (`--extent level`). ### Sources @@ -162,8 +165,8 @@ The first release, an alpha: the plugin API can still change before 0.1.0. level's whole shape, the box's chunks in it rounded out to whole chunks (`extent: "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; + (`computed`, `needed`); the Python side (`chunkmirage materialize --extent`) writes + either too, the box by default; `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/formats.md b/docs/concepts/formats.md index 52941a0..8119f00 100644 --- a/docs/concepts/formats.md +++ b/docs/concepts/formats.md @@ -90,10 +90,18 @@ 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, 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=, +region on its own voxels, rounded out to whole voxels. As the browser engine does, +`--extent box` (the default) writes exactly the box: each level an array of the box's +shape, its translation moved by the box's corner so it lands where the box is, in chunks +of the dataset's shape counted from that corner, those at the far edge stored whole with +zeros past the box. Each of its chunks is read from the dataset's chunks it overlaps; when +the last stage keeps no chunks of its own (no `cache: true`), those pass through a cache +of up to 1 GiB, walked so that the fewest wait (the axes the box's corner is on chunk +edges along first, then the one with the most chunks), so each is computed once while a +slab of them fits. `--extent level` writes the box's chunks, rounded out to whole chunks, +into arrays of the levels' whole shapes, for several boxes in one store; a box of the +whole level is the same store either way. From Python, +`chunkmirage.materialize(pipeline, where, format=, levels=, box=, extent=, 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 (`{chunks, written, skipped, bytes}`). Inside a running server it competes with clients' diff --git a/docs/reference/api.md b/docs/reference/api.md index 29eee42..6841df4 100644 --- a/docs/reference/api.md +++ b/docs/reference/api.md @@ -144,7 +144,7 @@ server.stop() address must be known before the app is built. `chunkmirage.materialize(pipeline, where, *, format="zarr3", levels=None, box=None, -compressor=None, skip_empty=True, progress=None)` writes a pipeline (or a box of it) into +extent="box", compressor=None, skip_empty=True, progress=None)` writes a pipeline (or a box of it) into a store as a frontend serves it, which reads back as a stored dataset: the Python side of `chunkmirage materialize` ([formats](../concepts/formats.md#a-dataset-written-to-a-store-python-and-the-command-line)). diff --git a/docs/reference/cli.md b/docs/reference/cli.md index d9a3d7c..f9c436d 100644 --- a/docs/reference/cli.md +++ b/docs/reference/cli.md @@ -66,7 +66,8 @@ 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, in an array of the level's whole shape (the browser engine's exact box, `extent: "box"`, is not supported here yet) | +| `--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 on its own voxels, rounded out to whole voxels | +| `--extent` | `box` | `box`: each level an array of exactly the box, placed where it is (its translation moved by the box's corner), in chunks counted from that corner; `level`: the box's chunks, rounded out to whole chunks, in an array of the level's whole shape (several boxes in one store) | | `--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/src/chunkmirage/cli.py b/src/chunkmirage/cli.py index 24a4f19..40be920 100644 --- a/src/chunkmirage/cli.py +++ b/src/chunkmirage/cli.py @@ -409,6 +409,10 @@ def materialize( help="the region, in level 0's voxels along the last axes: lo:hi, each comma separated in " "C order, e.g. 0,0,0:64,512,512 (every level writes the same region; default all of it)", ), + extent: str = typer.Option( + "box", help="box: an array of exactly the box, placed where it is; level: its chunks in an " + "array of the level's whole shape (to write several boxes into one store)", + ), chunk: str | None = typer.Option(None, help="output chunk shape, as for serve"), select: str | None = typer.Option(None, help="pin non-spatial axes, as for serve"), compressor: str | None = typer.Option(None, help="as the format takes it (default: blosc for zarr, gzip for n5)"), @@ -438,7 +442,7 @@ def report(c): counts = write( registry.get("processed"), out, format=format, levels=level or None, box=region, - compressor=compressor, skip_empty=not keep_empty, progress=report, + extent=extent, compressor=compressor, skip_empty=not keep_empty, progress=report, ) typer.echo(json.dumps({**counts, "seconds": round(time.perf_counter() - started, 1)})) diff --git a/src/chunkmirage/materialize.py b/src/chunkmirage/materialize.py index 55164ee..0bd1582 100644 --- a/src/chunkmirage/materialize.py +++ b/src/chunkmirage/materialize.py @@ -11,6 +11,12 @@ (``demand.gather``), and one of nothing but zeros (the fill value) is left out unless asked for. The store is anything tensorstore writes: a directory, or a ``file://``, ``gs://`` or ``s3://`` URL. + +A box is written as exactly the box (``extent="box"``, the default): each level an array +of the box's shape on it, placed where the box is (its translation moved by the box's +corner), in chunks of the dataset's shape counted from that corner, each read from the +dataset's chunks it overlaps. ``extent="level"`` writes the box's chunks into arrays of the +levels' whole shapes instead, so several boxes can go into one store. """ from __future__ import annotations @@ -22,13 +28,19 @@ import numpy as np from chunkmirage import demand -from chunkmirage.core import Box +from chunkmirage.cache import LRUCache +from chunkmirage.core import ArrayInfo, Box from chunkmirage.frontends import get_frontend from chunkmirage.frontends.base import Metadata from chunkmirage.pipeline import Pipeline +from chunkmirage.sources.base import ChunkedSource FORMATS = ("zarr3", "zarr", "n5") +EXTENTS = ("box", "level") BATCH = 256 # chunks asked for at once: the pool's work, and what one wait holds +#: most a box's write keeps of the dataset's chunks it stitches from, when its last stage +#: keeps none itself: enough for a slab of them, each computed once; past it, recomputed +BOX_BUFFER_BYTES = 1 << 30 def _kvstore_url(where: str) -> str: @@ -53,6 +65,50 @@ def level_box(pipeline: Pipeline, level: int, box: Box | None) -> Box: return Box((0,) * lead + tuple(lo), info.shape[:lead] + tuple(hi)).clip(info.shape) +class _Cropped: + """``pipeline`` cut to ``boxes``, one per level, as ``extent="box"`` writes it: each + level an array of its box's shape whose voxel 0 is the box's first (its translation + moved there), chunked as the level is, counted from the box's corner. A chunk is read + from the level's chunks it overlaps, through a cache of them when the level keeps none + (an uncached last stage), so each is computed once while a slab of them fits.""" + + def __init__(self, pipeline: Pipeline, boxes: list[Box]): + self.pipeline, self.source, self.boxes = pipeline, pipeline.source, boxes + self._infos: list[ArrayInfo] = [] + self._reads = [] + for level, b in enumerate(boxes): + info = pipeline.info(level) + moved = tuple(t + s * v for t, s, v in zip(info.translation, b.start, info.voxel_size)) + self._infos.append(info.with_(shape=b.shape, translation=moved)) + stage = pipeline.levels[level] + if not isinstance(stage, ChunkedSource) or stage.cached: + self._reads.append(stage) + else: + self._reads.append(ChunkedSource(stage.info, stage.chunk, LRUCache(BOX_BUFFER_BYTES), f"box:{stage.cache_key()}")) + + @property + def num_levels(self) -> int: + return self.pipeline.num_levels + + def info(self, level: int = 0) -> ArrayInfo: + return self._infos[level] + + def chunk(self, level: int, index) -> np.ndarray: + at, b = self._infos[level].chunk_box(index), self.boxes[level] + return self._reads[level].read(Box(tuple(s + a for s, a in zip(b.start, at.start)), + tuple(s + a for s, a in zip(b.start, at.stop)))) + + def walk(self, level: int) -> list[tuple[int, ...]]: + """The box's chunks in the order that keeps fewest of the level's open: the axes + its corner is on the level's chunk edges along first, then those it is not, the one + with the most chunks outermost (a chunk across an edge waits for the far side's).""" + info, b = self._infos[level], self.boxes[level] + grid = info.chunk_grid + split = [a for a in range(info.ndim) if b.start[a] % info.chunk_shape[a]] + order = [a for a in range(info.ndim) if a not in split] + sorted(split, key=lambda a: -grid[a]) + return sorted(np.ndindex(*grid), key=lambda idx: tuple(idx[a] for a in order)) + + def materialize( pipeline: Pipeline, where: str, @@ -60,6 +116,7 @@ def materialize( format: str = "zarr3", levels: Sequence[int] | None = None, box: Box | None = None, + extent: str = "box", compressor: str | None = None, skip_empty: bool = True, progress: Callable[[dict], None] | None = None, @@ -69,23 +126,34 @@ def materialize( frontend takes it (blosc by default for zarr, gzip for N5). ``levels``: which levels' chunks to write (all by default; every level's metadata is written whatever it is). ``box``: voxels of level 0 along its last axes, each level the same region on its own - voxels, rounded out to whole chunks. ``progress`` is told the counts after each batch - of chunks. Returns ``{chunks, written, skipped, bytes}``.""" + voxels (rounded out to whole voxels), written as ``extent`` says: ``box`` (the default) + an array of exactly the box, placed where it is; ``level`` its chunks (rounded out to + whole chunks) in an array of the level's whole shape. ``progress`` is told the counts + after each batch of chunks. Returns ``{chunks, written, skipped, bytes}``.""" import tensorstore as ts if format not in FORMATS: raise ValueError(f"format {format!r}: one of {', '.join(FORMATS)}") + if extent not in EXTENTS: + raise ValueError(f"extent {extent!r}: one of {', '.join(EXTENTS)}") fe = get_frontend(format, **({} if compressor is None else {"compressor": compressor})) kv = ts.KvStore.open(_kvstore_url(where)).result() levels = list(range(pipeline.num_levels)) if levels is None else [int(lv) for lv in levels] for lv in levels: if not 0 <= lv < pipeline.num_levels: raise ValueError(f"no level {lv}: the dataset has {pipeline.num_levels}") + source = pipeline + if box is not None and extent == "box": + boxes = [level_box(pipeline, lv, box) for lv in range(pipeline.num_levels)] + for lv, b in enumerate(boxes): + if b.empty: + raise ValueError(f"the box {box} holds no voxel of level {lv}, which is {pipeline.info(lv).shape}") + source = _Cropped(pipeline, boxes) # the metadata first: from then on the store is the dataset, zero where nothing is yet keys = [*fe.root_keys, *(f"s{i}/{k}" for i in range(pipeline.num_levels) for k in fe.level_keys)] for key in keys: - meta = fe.resolve(pipeline, key) + meta = fe.resolve(source, key) if isinstance(meta, Metadata): kv.write(key, meta.body).result() @@ -93,16 +161,18 @@ def materialize( def one(level: int, index: tuple[int, ...]) -> int | None: """Chunk ``index`` of ``level`` computed and written: its size, or None if empty.""" - block = pipeline.chunk(level, index) + block = source.chunk(level, index) if skip_empty and not block.any(): return None - data = fe.encode(pipeline.info(level), index, block) + data = fe.encode(source.info(level), index, block) kv.write(fe.chunk_key(level, index), data).result() return len(data) for level in levels: - info = pipeline.info(level) - todo = list(info.chunks_covering(level_box(pipeline, level, box))) + if isinstance(source, _Cropped): + todo = source.walk(level) + else: + todo = list(pipeline.info(level).chunks_covering(level_box(pipeline, level, box))) counts["chunks"] += len(todo) for start in range(0, len(todo), BATCH): for size in demand.gather([partial(one, level, idx) for idx in todo[start : start + BATCH]]): diff --git a/tests/test_materialize.py b/tests/test_materialize.py index fd5b267..b210ffe 100644 --- a/tests/test_materialize.py +++ b/tests/test_materialize.py @@ -1,6 +1,7 @@ """A dataset written to a store as the server serves it, read back as a stored dataset.""" import json +from typing import ClassVar import numpy as np import pytest @@ -9,6 +10,8 @@ from chunkmirage import Pipeline, PipelineSpec, materialize, open_source from chunkmirage.cli import app from chunkmirage.core import Box +from chunkmirage.ops import base +from chunkmirage.ops.base import Op IMG = "synthetic://julia?shape=32,48,48&chunk=16,16,16&levels=2&seed=2&voxel_size=1" # moved 24 voxels along x: the last third of every row reads past the image, zeros @@ -35,7 +38,7 @@ def test_a_dataset_written_reads_back_as_it_was_served(tmp_path, fmt): def test_a_box_of_one_level_from_the_command_line(tmp_path): out = tmp_path / "box.zarr" - result = CliRunner().invoke(app, ["materialize", IMG, str(out), "--level", "1", "--box", "0,0,0:16,16,16"]) + result = CliRunner().invoke(app, ["materialize", IMG, str(out), "--level", "1", "--box", "0,0,0:16,16,16", "--extent", "level"]) assert result.exit_code == 0, result.output assert json.loads(result.stdout.strip().splitlines()[-1])["chunks"] == 1 # level 1's voxels 0-8: one chunk back = open_source(str(out)).levels @@ -43,3 +46,41 @@ def test_a_box_of_one_level_from_the_command_line(tmp_path): got = back[1].read(Box((0, 0, 0), back[1].info.shape)) np.testing.assert_array_equal(got[:16, :16, :16], _whole(p, 1)[:16, :16, :16]) assert not got[16:].any() and not got[:, 16:].any() and not back[0].read(Box((0, 0, 0), (16, 16, 16))).any() + + +class Count(Op): + """Each block as it is, counted: how often the dataset's chunks are computed.""" + + name: ClassVar[str] = "count_test" + calls: ClassVar[int] = 0 + + def apply(self, block): + Count.calls += 1 + return block + + +@pytest.mark.parametrize("fmt", ["zarr3", "zarr"]) +def test_a_box_is_written_as_exactly_the_box_where_it_is(tmp_path, monkeypatch, fmt): + monkeypatch.setitem(base._REGISTRY, Count.name, Count) + p = Pipeline.from_spec(PipelineSpec(source=IMG, ops=[{"op": "count_test"}])) # its last stage keeps nothing + box = Box((3, 5, 7), (29, 40, 45)) # off the chunk edges along every axis + Count.calls = 0 + materialize(p, str(tmp_path / "box"), format=fmt, box=box) + # each of the dataset's chunks the box reaches computed once, though the box's own + # chunks each stitch up to eight of them + reached = sum(len(list(p.info(lv).chunks_covering(_level(p, lv, box)))) for lv in range(p.num_levels)) + assert Count.calls == reached + back = open_source(str(tmp_path / "box")).levels + got, want = back[0].info, p.info(0) + assert got.shape == box.shape and got.chunk_shape == want.chunk_shape + assert got.translation == tuple(t + s * v for t, s, v in zip(want.translation, box.start, want.voxel_size)) + np.testing.assert_array_equal(back[0].read(Box((0, 0, 0), box.shape)), p.read(0, box)) + lv1 = _level(p, 1, box) # level 1: the same region on its voxels, rounded out to whole ones + assert back[1].info.shape == lv1.shape + np.testing.assert_array_equal(back[1].read(Box((0, 0, 0), lv1.shape)), p.read(1, lv1)) + + +def _level(p, level, box): + from chunkmirage.materialize import level_box + + return level_box(p, level, box)