Spacelab 0.28.0__tar.gz

spacelab-0.28.0/.gitignore

@@ -0,0 +1,38 @@
+# Dependencies
+node_modules/
+
+# Build output
+dist/
+
+# Test output
+coverage/
+playwright-report/
+test-results/
+
+# Local env
+.env
+.env.local
+
+# Editor / OS cruft
+.DS_Store
+._*
+.last-review.md
+
+# Python
+__pycache__/
+*.py[cod]
+.ipynb_checkpoints/
+.venv/
+*.egg-info/
+
+# Personal Claude Code settings (per-clone)
+.claude/settings.local.json
+
+# Sphinx docs build
+docs/_build/
+
+# PyPI package build output (frontend copied in by py/build.sh)
+py/src/spacelab/static/
+py/src/spacelab/prep_spatialdata.py
+py/dist/
+py/LICENSE

spacelab-0.28.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 Allen Institute
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

spacelab-0.28.0/PKG-INFO

@@ -0,0 +1,30 @@
+Metadata-Version: 2.5
+Name: Spacelab
+Version: 0.28.0
+Summary: Browser-based viewer for bioimaging data (spatial transcriptomics, multiplexed immunofluorescence).
+Project-URL: Repository, https://github.com/aifimmunology/Spacelab
+License-Expression: MIT
+License-File: LICENSE
+Keywords: spatial-transcriptomics,visualization,xenium,zarr
+Requires-Python: >=3.10
+Provides-Extra: hlod
+Requires-Dist: numpy>=1.26; extra == 'hlod'
+Requires-Dist: pyarrow>=14; extra == 'hlod'
+Requires-Dist: zarr>=3.1; (python_version >= '3.11') and extra == 'hlod'
+Description-Content-Type: text/markdown
+
+# spacelab
+
+Browser-based viewer for spatial transcriptomics data (Xenium, Atera, CellDIVE).
+
+```bash
+uvx spacelab /path/to/data_app_ready   # or: pip install spacelab && spacelab
+```
+
+Run `spacelab` with no argument to serve the current directory. The directory must
+contain a `tissue-set-manifest.json` (the output of the conversion pipeline).
+
+Options: `--port 3000`, `--host 127.0.0.1`, `--no-browser`.
+
+See the [main repository](https://github.com/aifimmunology/Spacelab) for
+converting data.

spacelab-0.28.0/README.md

@@ -0,0 +1,15 @@
+# spacelab
+
+Browser-based viewer for spatial transcriptomics data (Xenium, Atera, CellDIVE).
+
+```bash
+uvx spacelab /path/to/data_app_ready   # or: pip install spacelab && spacelab
+```
+
+Run `spacelab` with no argument to serve the current directory. The directory must
+contain a `tissue-set-manifest.json` (the output of the conversion pipeline).
+
+Options: `--port 3000`, `--host 127.0.0.1`, `--no-browser`.
+
+See the [main repository](https://github.com/aifimmunology/Spacelab) for
+converting data.

spacelab-0.28.0/build.sh

@@ -0,0 +1,15 @@
+#!/usr/bin/env bash
+# Build the frontend, bundle it into the Python package, and build the wheel + sdist.
+set -euo pipefail
+here="$(cd "$(dirname "$0")" && pwd)"
+cd "$here/.."
+bun install --frozen-lockfile
+bun run build
+rm -rf "$here/src/spacelab/static"
+# Source maps are ~most of dist/ and only useful for debugging.
+rsync -a --exclude '*.map' dist/ "$here/src/spacelab/static/"
+cp LICENSE "$here/LICENSE"
+cp data/prep_spatialdata.py "$here/src/spacelab/prep_spatialdata.py"
+cd "$here"
+rm -rf dist
+uv build

spacelab-0.28.0/pyproject.toml

@@ -0,0 +1,40 @@
+[project]
+name = "Spacelab"
+version = "0.28.0"
+description = "Browser-based viewer for bioimaging data (spatial transcriptomics, multiplexed immunofluorescence)."
+readme = "README.md"
+license = "MIT"
+requires-python = ">=3.10"
+# The viewer is standard-library only; the frontend ships prebuilt in spacelab/static.
+dependencies = []
+keywords = ["spatial-transcriptomics", "xenium", "visualization", "zarr"]
+
+[project.optional-dependencies]
+# `spacelab hlod`: build/read/verify the Hilbert LOD transcript store (points.zarr).
+# zarr 3 needs Python 3.11+; the viewer itself stays 3.10-compatible.
+hlod = ["numpy>=1.26", "pyarrow>=14", "zarr>=3.1; python_version >= '3.11'"]
+
+[dependency-groups]
+dev = ["pytest>=8"]
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+
+[project.urls]
+Repository = "https://github.com/aifimmunology/Spacelab"
+
+[project.scripts]
+spacelab = "spacelab.cli:main"
+
+[build-system]
+requires = ["hatchling>=1.26"]
+build-backend = "hatchling.build"
+
+[tool.hatch.build.targets.wheel]
+packages = ["src/spacelab"]
+# static/ is gitignored build output (copied from the Vite dist/), so it must be
+# listed as an artifact or hatch will skip it.
+artifacts = ["src/spacelab/static/**", "src/spacelab/prep_spatialdata.py"]
+
+[tool.hatch.build.targets.sdist]
+artifacts = ["src/spacelab/static/**", "src/spacelab/prep_spatialdata.py"]

spacelab-0.28.0/src/spacelab/__init__.py

@@ -0,0 +1,6 @@
+from importlib.metadata import PackageNotFoundError, version
+
+try:
+    __version__ = version("spacelab")
+except PackageNotFoundError:  # imported from a source tree, not an install
+    __version__ = "0+unknown"

spacelab-0.28.0/src/spacelab/__main__.py

@@ -0,0 +1,3 @@
+from spacelab.cli import main
+
+main()

spacelab-0.28.0/src/spacelab/cli.py

@@ -0,0 +1,97 @@
+import argparse
+import os
+import shutil
+import subprocess
+import sys
+import threading
+import webbrowser
+from pathlib import Path
+
+from spacelab import __version__
+from spacelab.server import STATIC_DIR, make_server
+
+PREP_SCRIPT = Path(__file__).parent / "prep_spatialdata.py"
+
+MANIFEST = "tissue-set-manifest.json"
+
+
+def prep(args: list[str]) -> None:
+    if not PREP_SCRIPT.is_file():
+        sys.exit(f"error: prep script not bundled at {PREP_SCRIPT} (broken install?)")
+    # uv resolves the script's inline dependencies (zarr 3, shapely, ...) on the fly.
+    uv = shutil.which("uv")
+    cmd = [uv, "run", "--script", str(PREP_SCRIPT), *args] if uv else [sys.executable, str(PREP_SCRIPT), *args]
+    try:
+        sys.exit(subprocess.call(cmd))
+    except KeyboardInterrupt:
+        sys.exit(130)
+
+
+def hlod(args: list[str]) -> None:
+    try:
+        import numpy, pyarrow, zarr  # noqa: F401, E401
+    except ImportError:
+        # Same trick as prep: let uv supply the hlod extra's dependencies.
+        uv = shutil.which("uv")
+        if not uv:
+            sys.exit("error: `spacelab hlod` needs the hlod extra: pip install 'spacelab[hlod]'")
+        env = {**os.environ, "PYTHONPATH": str(Path(__file__).parent.parent)}
+        cmd = [uv, "run", "--no-project", "--with", "numpy>=1.26", "--with", "pyarrow>=14",
+               "--with", "zarr>=3.1", "python", "-m", "spacelab.hlod", *args]
+        try:
+            sys.exit(subprocess.call(cmd, env=env))
+        except KeyboardInterrupt:
+            sys.exit(130)
+    from spacelab.hlod.cli import main as hlod_main
+
+    sys.exit(hlod_main(args))
+
+
+def main(argv: list[str] | None = None) -> None:
+    argv = sys.argv[1:] if argv is None else argv
+    # Handled before argparse: the serve command's optional positional would swallow "prep".
+    if argv and argv[0] == "prep":
+        prep(argv[1:])
+        return
+    if argv and argv[0] == "hlod":
+        hlod(argv[1:])
+        return
+    parser = argparse.ArgumentParser(
+        prog="spacelab",
+        description="Serve a converted spatial dataset in the browser. "
+        "Run `spacelab prep --help` to add fast-loading sidecars to a SpatialData store, "
+        "or `spacelab hlod --help` for the transcript LOD store.",
+    )
+    parser.add_argument(
+        "data_dir", nargs="?", default=".", help="*_app_ready directory (default: cwd)"
+    )
+    parser.add_argument("--port", type=int, default=3000)
+    parser.add_argument("--host", default="127.0.0.1")
+    parser.add_argument("--no-browser", action="store_true", help="don't open a browser tab")
+    parser.add_argument("--version", action="version", version=f"spacelab {__version__}")
+    args = parser.parse_args(argv)
+
+    data_dir = Path(args.data_dir).expanduser().resolve()
+    # Without a manifest the app 404s on it and falls back to its connect-a-bucket page.
+    has_manifest = (data_dir / MANIFEST).is_file()
+    if not has_manifest:
+        print(f"No {MANIFEST} found in {data_dir}, so defaulting to cloud connect for data.")
+    if not (STATIC_DIR / "index.html").is_file():
+        sys.exit(f"error: frontend not bundled at {STATIC_DIR} (broken install?)")
+
+    try:
+        server = make_server(str(data_dir), args.host, args.port)
+    except OSError as e:
+        sys.exit(f"error: cannot bind {args.host}:{args.port} ({e}); try --port")
+
+    url = f"http://{'localhost' if args.host in ('127.0.0.1', '0.0.0.0') else args.host}:{args.port}"
+    target = f"{data_dir} " if has_manifest else ""
+    print(f"🚀 spacelab v{__version__} launching {target}on {url} (Ctrl+C to stop)", flush=True)
+    if not args.no_browser:
+        threading.Timer(0.5, webbrowser.open, [url]).start()
+    try:
+        server.serve_forever()
+    except KeyboardInterrupt:
+        print("\nstopped")
+    finally:
+        server.server_close()

spacelab-0.28.0/src/spacelab/hlod/__init__.py

@@ -0,0 +1,8 @@
+"""Hilbert additive-LOD transcript store: build, read and verify.
+
+Needs the `hlod` extra (numpy, pyarrow, zarr>=3): pip install 'spacelab[hlod]'.
+"""
+
+from .read import Gene, Store, open
+
+__all__ = ["Gene", "Store", "open"]

spacelab-0.28.0/src/spacelab/hlod/__main__.py

@@ -0,0 +1,5 @@
+import sys
+
+from .cli import main
+
+sys.exit(main())

spacelab-0.28.0/src/spacelab/hlod/build.py

@@ -0,0 +1,349 @@
+"""Build a Hilbert additive-LOD point store from a transcript table.
+
+Rows are sorted by (gene, level, h): a gene is one contiguous slice, levels
+0..d of a gene are one contiguous prefix of it, and within a level every
+quadtree cell at the level's index depth is one contiguous run. Levels are
+disjoint random subsets chosen by hashing a stable per-row key, so levels
+0..d together are a uniform random sample of fraction F[d].
+"""
+
+import math
+import os
+import shutil
+import tempfile
+import time
+from collections.abc import Callable, Sequence
+
+import numpy as np
+import pyarrow as pa
+import pyarrow.parquet as pq
+import zarr
+from zarr.codecs import ZstdCodec
+
+from .cells import block_offsets
+from .hashing import gene_xy_key, splitmix64
+from .hilbert import quantize, xy2d
+
+# Written to the root attrs. The browser reader (src/hlod/store.ts) refuses a
+# store with any other format or version, read.py checks the format, and
+# `hlod verify` the version; bump the version on any layout change.
+FORMAT = "hilbert-lod-points"
+FORMAT_VERSION = 1
+# Cap on the deepest level, Dmax. Each level is 4x the previous, so 6 levels past
+# level 0 reach P * 4^6 (~41M rows at the default P); a larger gene puts the rest
+# in its last level. Bounds the per-gene cell table, 4^(d+g) entries per level d.
+MAX_DMAX = 6
+# Array names the builder writes itself; an extra column with one of these names
+# would overwrite it.
+RESERVED = {"coords", "index", "original_index"}
+
+DEFAULTS = dict(
+    # Level-0 rows per gene: the subsample drawn when the whole tissue is in view.
+    # Smaller means sparser zoomed-out views and more levels.
+    P=10_000,
+    # Hilbert grid resolution: 2^K x 2^K over the square extent. 16 keeps h in uint32.
+    K=16,
+    # Extra quadtree depth for the cell index: level d is indexed at depth d + g,
+    # so level 0 already splits into 4^g cells for viewport queries.
+    g=1,
+    # Seed for the per-row hash that assigns rows to levels; fixed so rebuilds of
+    # the same input give the same store.
+    seed=42,
+    # Rows per zarr chunk, the unit of compression and of a byte-range read.
+    chunk=16_384,
+    # Rows per shard file (a multiple of chunk); fewer, larger files to serve.
+    shard=4_194_304,
+    # Where each shard file keeps its chunk index ("start" or "end", zarr's
+    # default); recorded in the array metadata, and the reader handles both.
+    index_location="end",
+)
+
+
+class BuildError(Exception):
+    """Bad input or options; the message is meant for the user."""
+
+
+def choose_dmax(gene_counts: np.ndarray, P: int) -> int:
+    """Deepest level any gene needs for its last level to hold all its rows,
+    in 1..MAX_DMAX."""
+    deepest = 0
+    for n in gene_counts:
+        if n > P:
+            deepest = max(deepest, math.ceil(math.log(n / P, 4)))
+    return max(1, min(MAX_DMAX, deepest))
+
+
+def level_thresholds(n_gene: int, P: int, dmax: int) -> list[int]:
+    """T[d] for d < dmax as uint64 cut points: a row is in level d when
+    T[d-1] <= r < T[d]. Integer comparison avoids float rounding at the edges."""
+    out = []
+    for d in range(dmax):
+        frac = min(1.0, P * 4 ** d / n_gene) if n_gene else 1.0
+        out.append(min(math.floor(frac * 2 ** 64), 2 ** 64 - 1))
+    return out
+
+
+def assign_levels(r: np.ndarray, n_gene: int, P: int, dmax: int) -> np.ndarray:
+    """Level (uint8) of each row of one gene, from its uint64 hash `r`."""
+    t = level_thresholds(n_gene, P, dmax)
+    level = np.searchsorted(np.array(t, dtype=np.uint64), r, side="right").astype(np.uint8)
+    # A level whose cumulative fraction is 1 takes every remaining row, including
+    # the one hash value (2^64 - 1) the clamped threshold would push past it.
+    full = next((d for d in range(dmax) if n_gene and P * 4 ** d >= n_gene), None)
+    if full is not None:
+        np.minimum(level, full, out=level)
+    return level
+
+
+def _read_table(source, cols: list[str]) -> pa.Table:
+    """`cols` of a pyarrow Table or parquet file/directory; BuildError if any is missing."""
+    if isinstance(source, pa.Table):
+        missing = [c for c in cols if c not in source.column_names]
+        if missing:
+            raise BuildError(f"input has no column(s) {missing}")
+        return source.select(cols)
+    schema_names = set(pq.read_schema(_first_parquet(source)).names)
+    missing = [c for c in cols if c not in schema_names]
+    if missing:
+        raise BuildError(f"{source} has no column(s) {missing} (has: {sorted(schema_names)})")
+    return pq.read_table(source, columns=cols)
+
+
+def _first_parquet(path: str) -> str:
+    """`path` itself if a file, else its first .parquet part (for reading the schema)."""
+    if os.path.isfile(path):
+        return path
+    parts = sorted(f for f in os.listdir(path) if f.endswith(".parquet"))
+    if not parts:
+        raise BuildError(f"{path} has no parquet files")
+    return os.path.join(path, parts[0])
+
+
+def _dictionary(column: pa.ChunkedArray) -> tuple[np.ndarray, list[str]]:
+    """(codes, names) with names sorted, so a code is the position in `names`.
+    Each parquet part carries its own dictionary, so they are unified first.
+    A null becomes the empty string."""
+    column = column.cast(pa.string()) if not pa.types.is_dictionary(column.type) else column
+    if not pa.types.is_dictionary(column.type):
+        column = column.dictionary_encode(null_encoding="encode")
+    arr = column.unify_dictionaries().combine_chunks()
+    if arr.null_count:
+        arr = pa.DictionaryArray.from_arrays(
+            arr.indices.fill_null(len(arr.dictionary)),
+            pa.concat_arrays([arr.dictionary.cast(pa.string()), pa.array([""])]))
+    raw = ["" if v is None else str(v) for v in arr.dictionary.to_pylist()]
+    order = sorted(range(len(raw)), key=lambda i: raw[i])
+    names, remap = [], np.empty(len(raw), dtype=np.int64)
+    for i in order:  # merge duplicate names (a unified dictionary can repeat after casting)
+        if not names or names[-1] != raw[i]:
+            names.append(raw[i])
+        remap[i] = len(names) - 1
+    codes = remap[arr.indices.to_numpy(zero_copy_only=False).astype(np.int64)]
+    return codes, names
+
+
+def _numeric(column: pa.ChunkedArray, dtype=None) -> np.ndarray:
+    """Column as a numpy array, float nulls as NaN, cast to `dtype` if given."""
+    arr = column.combine_chunks()
+    if arr.null_count and pa.types.is_floating(arr.type):
+        arr = arr.fill_null(float("nan"))
+    out = arr.to_numpy(zero_copy_only=False)
+    return out.astype(dtype) if dtype is not None else out
+
+
+def _is_categorical(t: pa.DataType) -> bool:
+    """String-like columns, stored as codes plus a `<column>_categories` array."""
+    return pa.types.is_string(t) or pa.types.is_large_string(t) or pa.types.is_dictionary(t)
+
+
+def build(source, output: str, *, gene_col: str, x_col: str, y_col: str,
+          id_col: str | None = None, columns: Sequence[str] = (),
+          P: int = DEFAULTS["P"], K: int = DEFAULTS["K"], g: int = DEFAULTS["g"],
+          seed: int = DEFAULTS["seed"], chunk: int = DEFAULTS["chunk"],
+          shard: int = DEFAULTS["shard"], index_location: str = DEFAULTS["index_location"],
+          original_index: bool = True, tmpdir: str | None = None,
+          extra_attrs: dict | None = None, log: Callable[[str], None] = print) -> dict:
+    """Write the store to `output` (replaced atomically via `output.tmp`).
+    `source` is a parquet file, a directory of parquet parts, or a pyarrow Table.
+    Returns a summary dict (also logged)."""
+    if shard % chunk:
+        raise BuildError(f"shard ({shard}) must be a multiple of chunk ({chunk})")
+    if index_location not in ("start", "end"):
+        raise BuildError("index_location must be 'start' or 'end'")
+    if not 1 <= K <= 16:
+        raise BuildError("K must be in 1..16 (h is stored in uint32)")
+    columns = list(columns)
+    bad = RESERVED.intersection(columns)
+    if bad:
+        raise BuildError(f"column name(s) {sorted(bad)} are reserved")
+    t0 = time.time()
+
+    read_cols = list(dict.fromkeys([gene_col, x_col, y_col, *([id_col] if id_col else []),
+                                    *columns]))
+    tbl = _read_table(source, read_cols)
+    n = tbl.num_rows
+    gene_codes, genes = _dictionary(tbl.column(gene_col))
+    # Quantize the float32 values actually stored, so verify (which only sees the
+    # stored coords) recomputes exactly the h the build sorted by.
+    x = _numeric(tbl.column(x_col), np.float32)
+    y = _numeric(tbl.column(y_col), np.float32)
+    if n and not (np.isfinite(x).all() and np.isfinite(y).all()):
+        raise BuildError("x/y contain NaN or infinite values")
+    ids = _numeric(tbl.column(id_col)).astype(np.uint64) if id_col else None
+
+    extras: dict[str, np.ndarray] = {}
+    categories: dict[str, list[str]] = {}
+    for c in columns:
+        col = tbl.column(c)
+        if _is_categorical(col.type):
+            codes, names = _dictionary(col)
+            extras[c] = codes.astype(np.uint32)
+            categories[c] = names
+        else:
+            extras[c] = _numeric(col)
+    del tbl
+
+    x0 = float(x.min()) if n else 0.0
+    y0 = float(y.min()) if n else 0.0
+    span = max(float(x.max()) - x0, float(y.max()) - y0) if n else 0.0
+    size = span * (1 + 1e-9) if span > 0 else 1.0
+
+    G = len(genes)
+    counts = np.bincount(gene_codes, minlength=G).astype(np.int64)
+    dmax = choose_dmax(counts, P)
+    if dmax + g > K:
+        raise BuildError(f"index depth {dmax + g} exceeds K={K}")
+    B = block_offsets(g, dmax)
+    S = int(B[-1])
+    gene_offsets = np.concatenate([[0], np.cumsum(counts)]).astype(np.uint64)
+    order_g = np.argsort(gene_codes, kind="stable")
+    del gene_codes
+
+    tmp_root = tempfile.mkdtemp(prefix="hlod-", dir=tmpdir)
+    try:
+        def mm(name, dtype, shape):
+            return np.lib.format.open_memmap(os.path.join(tmp_root, f"{name}.npy"), mode="w+",
+                                             dtype=dtype, shape=shape)
+
+        out_coords = mm("coords", np.float32, (n, 2))
+        out_orig = mm("original_index", np.uint64, (n,)) if original_index else None
+        out_extra = {c: mm(f"col_{i}", a.dtype, (n,)) for i, (c, a) in enumerate(extras.items())}
+        level_offsets = np.zeros((G, dmax + 2), dtype=np.uint64)
+        cell_index = np.zeros((G, S), dtype=np.uint32)
+        level_rows = np.zeros((G, dmax + 1), dtype=np.int64)
+
+        for gi in range(G):
+            a, b = int(gene_offsets[gi]), int(gene_offsets[gi + 1])
+            idx = order_g[a:b]
+            gx, gy = x[idx], y[idx]
+            h = xy2d(K, quantize(gx, x0, size, K), quantize(gy, y0, size, K))
+            key = ids[idx] if ids is not None else gene_xy_key(gi, gx, gy)
+            r = splitmix64(key ^ np.uint64(seed))
+            level = assign_levels(r, b - a, P, dmax)
+            order = np.lexsort((h, level))
+            h, level, idx = h[order], level[order], idx[order]
+
+            out_coords[a:b, 0] = gx[order]
+            out_coords[a:b, 1] = gy[order]
+            if out_orig is not None:
+                out_orig[a:b] = idx
+            for c, arr in extras.items():
+                out_extra[c][a:b] = arr[idx]
+
+            per_level = np.bincount(level, minlength=dmax + 1)
+            level_rows[gi] = per_level
+            rel = np.concatenate([[0], np.cumsum(per_level)])
+            level_offsets[gi] = a + rel
+            for d in range(dmax + 1):
+                m = d + g
+                cells = (h[rel[d]:rel[d + 1]] >> np.uint32(2 * (K - m))).astype(np.int64)
+                hist = np.bincount(cells, minlength=4 ** m)
+                cell_index[gi, B[d]:B[d + 1]] = rel[d] + np.concatenate([[0], np.cumsum(hist)])
+
+        attrs = {
+            "format": FORMAT,
+            "format_version": FORMAT_VERSION,
+            "K": K, "g": g, "P": P, "Dmax": dmax,
+            "extent": {"x0": x0, "y0": y0, "size": size},
+            "hash": {"algo": "splitmix64", "seed": seed,
+                     "key": "transcript_id" if id_col else "gene_xy"},
+            "genes": genes,
+            "columns": columns,
+            "categories": {c: f"{c}_categories" for c in categories},
+            "coordinate_units": "micron",
+            "chunk": chunk, "shard": shard, "index_location": index_location,
+            "source": {"path": source if isinstance(source, str) else None, "rows": n},
+            **(extra_attrs or {}),
+        }
+        _write_store(output, attrs, tmp_root, out_coords, out_orig, out_extra, categories,
+                     gene_offsets, level_offsets, cell_index, chunk, shard, index_location)
+    finally:
+        shutil.rmtree(tmp_root, ignore_errors=True)
+
+    index_bytes = S * 4 + (dmax + 2) * 8 + 8
+    largest = int(np.diff(cell_index.astype(np.int64), axis=1).max()) if G and S > 1 else 0
+    summary = {"path": output, "rows": n, "genes": G, "Dmax": dmax, "S": S,
+               "index_bytes_per_gene": index_bytes, "largest_cell_rows": largest,
+               "rows_per_level": level_rows.sum(axis=0).tolist(),
+               "seconds": round(time.time() - t0, 1)}
+    log(f"hlod: {n:,} rows, {G} genes, Dmax={dmax}, rows/level={summary['rows_per_level']}, "
+        f"index {index_bytes:,} B/gene, largest cell {largest:,} rows, {summary['seconds']}s")
+    return summary
+
+
+def _sharded(group, name, data: np.ndarray, chunk_rows: int, shard_rows: int,
+             index_location: str, compress: bool):
+    """Write `data` as a sharded zarr array, chunked along rows. Chunk and shard
+    sizes are clamped to the array."""
+    rows = data.shape[0]
+    chunk_rows = max(1, min(chunk_rows, rows))
+    # Whole multiple of the chunk; a shard larger than the array just holds every chunk.
+    shard_rows = max(chunk_rows, min(shard_rows, -(-rows // chunk_rows) * chunk_rows))
+    shard_rows -= shard_rows % chunk_rows
+    tail = data.shape[1:]
+    arr = group.create_array(
+        name, shape=data.shape, dtype=data.dtype,
+        chunks=(chunk_rows, *tail), shards={"shape": (shard_rows, *tail),
+                                            "index_location": index_location},
+        compressors=ZstdCodec(level=3) if compress else None,
+        # Keep every chunk present so uncompressed byte offsets stay contiguous.
+        config={"write_empty_chunks": True},
+    )
+    # Whole shards at a time: a partial shard write is a read-modify-write.
+    for s in range(0, rows, shard_rows):
+        arr[s:s + shard_rows] = np.asarray(data[s:s + shard_rows])
+    return arr
+
+
+def _write_store(output, attrs, tmp_root, coords, orig, extra, categories, gene_offsets,
+                 level_offsets, cell_index, chunk, shard, index_location):
+    """Write every array and the attrs into `output`.tmp, consolidate the metadata,
+    then move it to `output`."""
+    staging = output.rstrip("/") + ".tmp"
+    shutil.rmtree(staging, ignore_errors=True)
+    root = zarr.open_group(staging, mode="w", zarr_format=3)
+    _sharded(root, "coords", coords, chunk, shard, index_location, compress=False)
+    if orig is not None:
+        _sharded(root, "original_index", orig, chunk, shard, index_location, compress=True)
+    for c, arr in extra.items():
+        _sharded(root, c, arr, chunk, shard, index_location, compress=True)
+        if c in categories:
+            names = np.array(categories[c], dtype=object)
+            cat = root.create_array(f"{c}_categories", shape=names.shape, dtype=str,
+                                    chunks=(max(1, min(names.size, 1 << 20)),),
+                                    compressors=ZstdCodec(level=3))
+            if names.size:
+                cat[:] = names
+    index = root.create_group("index")
+    G, S = cell_index.shape
+    # Uncompressed so the client reads them with the same byte-range reader as coords.
+    _sharded(index, "gene_offsets", gene_offsets, gene_offsets.size, gene_offsets.size,
+             index_location, compress=False)
+    _sharded(index, "level_offsets", level_offsets, max(1, G), max(1, G), index_location,
+             compress=False)
+    _sharded(index, "cell_index", cell_index, 1, max(1, G), index_location, compress=False)
+    root.attrs.update(attrs)
+    zarr.consolidate_metadata(staging)
+    if os.path.exists(output):
+        shutil.rmtree(output)
+    os.replace(staging, output)