chunkmirage materialize writes exactly the box, as the browser engine does

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

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'

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)).

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) |

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)}))
 

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,13 +65,58 @@ 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,
     *,
     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,40 +126,53 @@ 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()
 
     counts = {"chunks": 0, "written": 0, "skipped": 0, "bytes": 0}
 
     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]]):

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,11 +38,49 @@ 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
     p = Pipeline.from_spec(PipelineSpec(source=IMG))
     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)