From 618b350708ce6e507f540907e8d3d79cb2d8f848 Mon Sep 17 00:00:00 2001 From: Nicolas Fernandez Date: Sun, 6 Sep 2026 13:43:03 -0400 Subject: [PATCH 01/17] feat(experimental): add deterministic regular-grid tiling Adds the viewer-independent core of the celldega_regular_grid_v1 profile: the tile geometry and the tile -> row-group -> file numbering. - x-major tile ids (tile_x * num_tiles_y + tile_y), matching Celldega's RowGroupTileReader, but reimplemented rather than imported so spatialdata-io gains no dependency on a viewer package (Celldega already depends on spatialdata-io, so importing it would be circular). A conformance test pins the two formulas together. - Half-open tile bounds, with the dataset's upper edge clamped into the last tile so points on x_max/y_max are not dropped. - Out-of-grid coordinates raise instead of being silently clamped, since they indicate a mismatched transform rather than a rounding artefact. - Zero-padded chunk filenames: Celldega indexes the manifest's files array by position while dask globs and sorts lexicographically, so unpadded names would make dask reorder partitions past 10 files. Co-Authored-By: Claude Opus 5 (1M context) --- .../experimental/regular_grid.py | 247 ++++++++++++++++++ tests/test_regular_grid.py | 233 +++++++++++++++++ 2 files changed, 480 insertions(+) create mode 100644 src/spatialdata_io/experimental/regular_grid.py create mode 100644 tests/test_regular_grid.py diff --git a/src/spatialdata_io/experimental/regular_grid.py b/src/spatialdata_io/experimental/regular_grid.py new file mode 100644 index 00000000..84a89b76 --- /dev/null +++ b/src/spatialdata_io/experimental/regular_grid.py @@ -0,0 +1,247 @@ +"""Deterministic regular-grid spatial tiling. + +This module defines the tile geometry and the tile -> row-group -> file numbering +used by the ``celldega_regular_grid_v1`` visualization profile. + +The grid is a non-overlapping regular square grid in *display pixel* space (level-0 +pixels of a chosen reference image). Given an origin, a tile size and grid dimensions, +a point's tile is a pure function of its coordinates:: + + tile_x = floor((x_px - origin_x) / tile_size_px) + tile_y = floor((y_px - origin_y) / tile_size_px) + tile_id = tile_x * num_tiles_y + tile_y + +Tile bounds are half-open ``[min, max)``, except at the upper edge of the dataset where +the last tile is closed so that points lying exactly on ``x_max`` / ``y_max`` are kept. + +``tile_id`` doubles as the *global row-group index*: one logical tile is written as +exactly one Parquet row group, including empty tiles, which are written as zero-row row +groups. Row groups are then split across files:: + + file_index = tile_id // max_row_groups_per_file + local_row_group = tile_id % max_row_groups_per_file + +Multi-file output is deliberate: a Parquet reader must fetch a file's entire footer +before it can read any row group, and footer size grows with row-group count. Splitting +keeps the browser's cold-start cost proportional to the viewport rather than to the +whole dataset. + +This numbering matches Celldega's ``RowGroupTileReader`` exactly, but nothing here is +Celldega-specific -- it is the viewer-independent part of the profile. +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import Any + +import numpy as np +from numpy.typing import NDArray + +__all__ = ["RegularGrid", "DEFAULT_MAX_ROW_GROUPS_PER_FILE"] + +#: Default number of row groups per Parquet file. Matches Celldega's writer default. +DEFAULT_MAX_ROW_GROUPS_PER_FILE = 400 + + +@dataclass(frozen=True) +class RegularGrid: + """A deterministic non-overlapping regular square grid in display-pixel space. + + Parameters + ---------- + origin_x + X coordinate (display pixels) of the grid origin, i.e. the left edge of column 0. + origin_y + Y coordinate (display pixels) of the grid origin, i.e. the top edge of row 0. + tile_size_px + Edge length of a tile, in display pixels. Must be positive. + num_tiles_x + Number of tile columns. Must be positive. + num_tiles_y + Number of tile rows. Must be positive. + """ + + origin_x: float + origin_y: float + tile_size_px: float + num_tiles_x: int + num_tiles_y: int + + def __post_init__(self) -> None: + if not self.tile_size_px > 0: + raise ValueError(f"tile_size_px must be positive, got {self.tile_size_px}") + if self.num_tiles_x < 1 or self.num_tiles_y < 1: + raise ValueError( + f"grid must have at least one tile per axis, got " + f"num_tiles_x={self.num_tiles_x}, num_tiles_y={self.num_tiles_y}" + ) + + # -- construction --------------------------------------------------------- + + @classmethod + def from_bounds( + cls, + x_min: float, + y_min: float, + x_max: float, + y_max: float, + tile_size_px: float, + ) -> RegularGrid: + """Build the smallest grid with the given tile size that covers ``[x_min, x_max] x [y_min, y_max]``. + + The origin is placed at ``(x_min, y_min)``. The grid always has at least one tile + per axis, so a degenerate (zero-extent) bounding box still yields a valid grid. + """ + if x_max < x_min or y_max < y_min: + raise ValueError(f"invalid bounds: ({x_min}, {y_min}) to ({x_max}, {y_max})") + num_tiles_x = max(1, math.ceil((x_max - x_min) / tile_size_px)) + num_tiles_y = max(1, math.ceil((y_max - y_min) / tile_size_px)) + return cls( + origin_x=float(x_min), + origin_y=float(y_min), + tile_size_px=float(tile_size_px), + num_tiles_x=int(num_tiles_x), + num_tiles_y=int(num_tiles_y), + ) + + # -- geometry ------------------------------------------------------------- + + @property + def num_tiles(self) -> int: + """Total number of logical tiles, which equals the total number of row groups.""" + return self.num_tiles_x * self.num_tiles_y + + @property + def x_max(self) -> float: + """Right edge of the grid in display pixels.""" + return self.origin_x + self.num_tiles_x * self.tile_size_px + + @property + def y_max(self) -> float: + """Bottom edge of the grid in display pixels.""" + return self.origin_y + self.num_tiles_y * self.tile_size_px + + def tile_xy( + self, x: NDArray[np.floating[Any]], y: NDArray[np.floating[Any]] + ) -> tuple[NDArray[np.int64], NDArray[np.int64]]: + """Map display-pixel coordinates to ``(tile_x, tile_y)`` indices. + + Coordinates on the upper edge of the grid are clamped into the last tile, so that + a point at exactly ``x_max`` belongs to tile column ``num_tiles_x - 1`` rather + than falling outside the grid. + + Raises + ------ + ValueError + If any coordinate falls below the grid origin, or more than one tile beyond + the grid extent. Such points indicate a mismatched grid rather than a + rounding artefact, and are not silently clamped. + """ + x = np.asarray(x) + y = np.asarray(y) + if x.shape != y.shape: + raise ValueError(f"x and y must have the same shape, got {x.shape} and {y.shape}") + + tx = np.floor((x - self.origin_x) / self.tile_size_px).astype(np.int64) + ty = np.floor((y - self.origin_y) / self.tile_size_px).astype(np.int64) + + self._check_in_range(tx, "x", self.num_tiles_x) + self._check_in_range(ty, "y", self.num_tiles_y) + + # Clamp the closed upper edge into the last tile. + np.clip(tx, 0, self.num_tiles_x - 1, out=tx) + np.clip(ty, 0, self.num_tiles_y - 1, out=ty) + return tx, ty + + def _check_in_range(self, t: NDArray[np.int64], axis: str, num_tiles: int) -> None: + if t.size == 0: + return + lo = int(t.min()) + hi = int(t.max()) + # `hi == num_tiles` is the closed upper edge and is clamped; anything beyond is an error. + if lo < 0 or hi > num_tiles: + raise ValueError( + f"{axis} coordinates map to tile indices [{lo}, {hi}], outside the grid " + f"[0, {num_tiles - 1}] (upper edge {num_tiles} allowed). " + f"The grid does not match the data; check origin, tile size and the " + f"micron-to-pixel transform." + ) + + def tile_id(self, tile_x: NDArray[np.int64], tile_y: NDArray[np.int64]) -> NDArray[np.int64]: + """Combine tile indices into a global tile id (x-major: ``tile_x * num_tiles_y + tile_y``).""" + return np.asarray(tile_x, dtype=np.int64) * self.num_tiles_y + np.asarray(tile_y, dtype=np.int64) + + def assign(self, x: NDArray[np.floating[Any]], y: NDArray[np.floating[Any]]) -> NDArray[np.int64]: + """Map display-pixel coordinates directly to global tile ids.""" + tx, ty = self.tile_xy(x, y) + return self.tile_id(tx, ty) + + def tile_bounds(self, tile_x: int, tile_y: int) -> tuple[float, float, float, float]: + """Return the half-open bounds ``(x_min, y_min, x_max, y_max)`` of one tile.""" + if not (0 <= tile_x < self.num_tiles_x and 0 <= tile_y < self.num_tiles_y): + raise ValueError(f"tile ({tile_x}, {tile_y}) is outside the grid") + x0 = self.origin_x + tile_x * self.tile_size_px + y0 = self.origin_y + tile_y * self.tile_size_px + return (x0, y0, x0 + self.tile_size_px, y0 + self.tile_size_px) + + # -- row-group / file numbering ------------------------------------------- + + def chunk_location( + self, tile_id: int, max_row_groups_per_file: int = DEFAULT_MAX_ROW_GROUPS_PER_FILE + ) -> tuple[int, int]: + """Return ``(file_index, local_row_group_index)`` for a global tile id.""" + return divmod(int(tile_id), int(max_row_groups_per_file)) + + def num_files(self, max_row_groups_per_file: int = DEFAULT_MAX_ROW_GROUPS_PER_FILE) -> int: + """Number of Parquet files needed to hold every tile's row group.""" + return math.ceil(self.num_tiles / max_row_groups_per_file) + + def chunk_filenames( + self, + max_row_groups_per_file: int = DEFAULT_MAX_ROW_GROUPS_PER_FILE, + prefix: str = "chunk", + ) -> list[str]: + """Return the ordered chunk filenames, zero-padded so lexicographic == numeric order. + + Padding matters because the two consumers disagree about ordering: Celldega indexes + the manifest's ``files`` array by position, while dask's ``read_parquet`` globs a + directory and sorts lexicographically. Without padding, ``chunk_10`` would sort + before ``chunk_2`` and dask would silently reorder partitions. + """ + n = self.num_files(max_row_groups_per_file) + width = len(str(n - 1)) if n > 1 else 1 + return [f"{prefix}_{i:0{width}d}.parquet" for i in range(n)] + + # -- serialization -------------------------------------------------------- + + def to_manifest_dict(self) -> dict[str, Any]: + """Serialize the grid into the ``tile_grid`` block of the profile manifest. + + Uses Celldega's established ``landscape_parameters.json`` key names so that an + existing reader can consume it unchanged. + """ + return { + "num_tiles_x": self.num_tiles_x, + "num_tiles_y": self.num_tiles_y, + "tile_size": self.tile_size_px, + "x_min": self.origin_x, + "y_min": self.origin_y, + "x_max": self.x_max, + "y_max": self.y_max, + } + + @classmethod + def from_manifest_dict(cls, d: dict[str, Any]) -> RegularGrid: + """Rebuild a grid from its manifest representation.""" + missing = {"num_tiles_x", "num_tiles_y", "tile_size", "x_min", "y_min"} - set(d) + if missing: + raise ValueError(f"tile_grid is missing required keys: {sorted(missing)}") + return cls( + origin_x=float(d["x_min"]), + origin_y=float(d["y_min"]), + tile_size_px=float(d["tile_size"]), + num_tiles_x=int(d["num_tiles_x"]), + num_tiles_y=int(d["num_tiles_y"]), + ) diff --git a/tests/test_regular_grid.py b/tests/test_regular_grid.py new file mode 100644 index 00000000..3cacf515 --- /dev/null +++ b/tests/test_regular_grid.py @@ -0,0 +1,233 @@ +"""Tests for the deterministic regular-grid tiling used by the visualization profile. + +These pin the *normative* parts of the profile: the tile formula, half-open boundary +semantics, x-major tile numbering, and the row-group/file numbering. A change that +breaks one of these invalidates every store already written with the profile. +""" + +from __future__ import annotations + +import numpy as np +import pytest + +from spatialdata_io.experimental.regular_grid import ( + DEFAULT_MAX_ROW_GROUPS_PER_FILE, + RegularGrid, +) + + +@pytest.fixture +def grid_2x3() -> RegularGrid: + """The synthetic fixture grid: 2 tile columns x 3 tile rows, 10 px tiles, origin at 0.""" + return RegularGrid(origin_x=0.0, origin_y=0.0, tile_size_px=10.0, num_tiles_x=2, num_tiles_y=3) + + +# -- construction ------------------------------------------------------------- + + +def test_from_bounds_covers_extent() -> None: + grid = RegularGrid.from_bounds(0, 0, 25000, 20000, tile_size_px=250) + assert (grid.num_tiles_x, grid.num_tiles_y) == (100, 80) + assert grid.num_tiles == 8000 + assert grid.x_max >= 25000 and grid.y_max >= 20000 + + +def test_from_bounds_rounds_up_partial_tiles() -> None: + # 105 px of data at 50 px tiles needs 3 columns, not 2. + grid = RegularGrid.from_bounds(0, 0, 105, 40, tile_size_px=50) + assert (grid.num_tiles_x, grid.num_tiles_y) == (3, 1) + + +def test_from_bounds_degenerate_extent_still_valid() -> None: + grid = RegularGrid.from_bounds(5, 5, 5, 5, tile_size_px=10) + assert grid.num_tiles == 1 + + +def test_from_bounds_nonzero_origin() -> None: + grid = RegularGrid.from_bounds(100, 200, 140, 260, tile_size_px=20) + assert (grid.origin_x, grid.origin_y) == (100.0, 200.0) + assert (grid.num_tiles_x, grid.num_tiles_y) == (2, 3) + # A point at the origin lands in tile (0, 0), not somewhere negative. + assert grid.assign(np.array([100.0]), np.array([200.0]))[0] == 0 + + +@pytest.mark.parametrize("bad", [0, -1]) +def test_rejects_nonpositive_tile_size(bad: float) -> None: + with pytest.raises(ValueError, match="tile_size_px must be positive"): + RegularGrid(0, 0, bad, 2, 3) + + +def test_rejects_empty_grid() -> None: + with pytest.raises(ValueError, match="at least one tile"): + RegularGrid(0, 0, 10, 0, 3) + + +def test_from_bounds_rejects_inverted_bounds() -> None: + with pytest.raises(ValueError, match="invalid bounds"): + RegularGrid.from_bounds(10, 0, 0, 10, tile_size_px=5) + + +# -- tile assignment ---------------------------------------------------------- + + +def test_tile_id_is_x_major(grid_2x3: RegularGrid) -> None: + """tile_id = tile_x * num_tiles_y + tile_y, matching Celldega's RowGroupTileReader.""" + expected = {(0, 0): 0, (0, 1): 1, (0, 2): 2, (1, 0): 3, (1, 1): 4, (1, 2): 5} + for (tx, ty), tid in expected.items(): + assert int(grid_2x3.tile_id(np.array(tx), np.array(ty))) == tid + + +def test_tile_ids_cover_every_tile_exactly_once(grid_2x3: RegularGrid) -> None: + ids = [ + int(grid_2x3.tile_id(np.array(tx), np.array(ty))) + for tx in range(grid_2x3.num_tiles_x) + for ty in range(grid_2x3.num_tiles_y) + ] + assert sorted(ids) == list(range(grid_2x3.num_tiles)) + + +def test_boundary_is_half_open(grid_2x3: RegularGrid) -> None: + """A point exactly on an internal tile edge belongs to the *upper* tile.""" + tx, ty = grid_2x3.tile_xy(np.array([9.999, 10.0]), np.array([0.0, 0.0])) + assert tx.tolist() == [0, 1] + assert ty.tolist() == [0, 0] + + +def test_upper_edge_is_clamped_into_last_tile(grid_2x3: RegularGrid) -> None: + """Points exactly on x_max / y_max stay in the grid rather than falling off it.""" + tx, ty = grid_2x3.tile_xy(np.array([20.0]), np.array([30.0])) + assert (int(tx[0]), int(ty[0])) == (1, 2) + assert int(grid_2x3.assign(np.array([20.0]), np.array([30.0]))[0]) == 5 + + +def test_out_of_range_raises_rather_than_clamping(grid_2x3: RegularGrid) -> None: + with pytest.raises(ValueError, match="outside the grid"): + grid_2x3.tile_xy(np.array([-1.0]), np.array([0.0])) + with pytest.raises(ValueError, match="outside the grid"): + grid_2x3.tile_xy(np.array([100.0]), np.array([0.0])) + + +def test_assign_is_vectorized_and_order_preserving(grid_2x3: RegularGrid) -> None: + x = np.array([0.0, 15.0, 5.0, 19.0]) + y = np.array([0.0, 25.0, 12.0, 5.0]) + got = grid_2x3.assign(x, y) + assert got.tolist() == [0, 5, 1, 3] + + +def test_assign_empty_input(grid_2x3: RegularGrid) -> None: + got = grid_2x3.assign(np.array([]), np.array([])) + assert got.shape == (0,) + + +def test_mismatched_shapes_raise(grid_2x3: RegularGrid) -> None: + with pytest.raises(ValueError, match="same shape"): + grid_2x3.tile_xy(np.array([1.0, 2.0]), np.array([1.0])) + + +def test_tile_bounds_are_contiguous(grid_2x3: RegularGrid) -> None: + _, _, x_max_0, _ = grid_2x3.tile_bounds(0, 0) + x_min_1, _, _, _ = grid_2x3.tile_bounds(1, 0) + assert x_max_0 == x_min_1 + + +def test_tile_bounds_rejects_out_of_grid(grid_2x3: RegularGrid) -> None: + with pytest.raises(ValueError, match="outside the grid"): + grid_2x3.tile_bounds(2, 0) + + +def test_every_tile_center_maps_back_to_its_own_tile(grid_2x3: RegularGrid) -> None: + """Round-trip: tile -> center point -> tile.""" + for tx in range(grid_2x3.num_tiles_x): + for ty in range(grid_2x3.num_tiles_y): + x0, y0, x1, y1 = grid_2x3.tile_bounds(tx, ty) + cx, cy = (x0 + x1) / 2, (y0 + y1) / 2 + got_x, got_y = grid_2x3.tile_xy(np.array([cx]), np.array([cy])) + assert (int(got_x[0]), int(got_y[0])) == (tx, ty) + + +# -- row-group and file numbering --------------------------------------------- + + +def test_chunk_location_formula() -> None: + grid = RegularGrid.from_bounds(0, 0, 25000, 20000, tile_size_px=250) + assert grid.chunk_location(0, 400) == (0, 0) + assert grid.chunk_location(399, 400) == (0, 399) + assert grid.chunk_location(400, 400) == (1, 0) + assert grid.chunk_location(7999, 400) == (19, 399) + + +def test_num_files_rounds_up() -> None: + grid = RegularGrid.from_bounds(0, 0, 100, 100, tile_size_px=10) # 100 tiles + assert grid.num_files(400) == 1 + assert grid.num_files(40) == 3 # 100 / 40 -> 3 files + assert grid.num_files(100) == 1 + + +def test_default_max_row_groups_matches_celldega() -> None: + assert DEFAULT_MAX_ROW_GROUPS_PER_FILE == 400 + + +def test_chunk_filenames_are_zero_padded_for_lexicographic_order() -> None: + """dask globs and sorts lexicographically; Celldega indexes by position. Padding satisfies both.""" + grid = RegularGrid.from_bounds(0, 0, 25000, 20000, tile_size_px=250) # 8000 tiles -> 20 files + names = grid.chunk_filenames(400) + assert len(names) == 20 + assert names[0] == "chunk_00.parquet" + assert names[10] == "chunk_10.parquet" + assert names == sorted(names), "lexicographic order must equal numeric order" + + +def test_chunk_filenames_single_file_unpadded() -> None: + grid = RegularGrid.from_bounds(0, 0, 100, 100, tile_size_px=50) # 4 tiles -> 1 file + assert grid.chunk_filenames(400) == ["chunk_0.parquet"] + + +def test_every_tile_maps_to_a_listed_file() -> None: + grid = RegularGrid.from_bounds(0, 0, 25000, 20000, tile_size_px=250) + names = grid.chunk_filenames(400) + for tid in range(grid.num_tiles): + file_index, local = grid.chunk_location(tid, 400) + assert 0 <= file_index < len(names) + assert 0 <= local < 400 + + +# -- manifest round-trip ------------------------------------------------------ + + +def test_manifest_round_trip(grid_2x3: RegularGrid) -> None: + assert RegularGrid.from_manifest_dict(grid_2x3.to_manifest_dict()) == grid_2x3 + + +def test_manifest_uses_celldega_key_names(grid_2x3: RegularGrid) -> None: + d = grid_2x3.to_manifest_dict() + assert set(d) == {"num_tiles_x", "num_tiles_y", "tile_size", "x_min", "y_min", "x_max", "y_max"} + + +def test_manifest_rejects_missing_keys() -> None: + with pytest.raises(ValueError, match="missing required keys"): + RegularGrid.from_manifest_dict({"num_tiles_x": 2, "num_tiles_y": 3}) + + +# -- conformance with Celldega's reader --------------------------------------- + + +def test_matches_celldega_row_group_index_formula() -> None: + """Mirror of RowGroupTileReader.computeRowGroupIndex / computeChunkLocation. + + Celldega is the reference client; this reimplements its JS formulas independently so + that a divergence in either codebase fails here rather than in the browser. + """ + grid = RegularGrid.from_bounds(0, 0, 25000, 20000, tile_size_px=250) + max_rg = 400 + rng = np.random.default_rng(0) + for tx, ty in zip( + rng.integers(0, grid.num_tiles_x, 200), + rng.integers(0, grid.num_tiles_y, 200), + strict=True, + ): + js_row_group = int(tx) * grid.num_tiles_y + int(ty) # computeRowGroupIndex + js_file = js_row_group // max_rg # computeChunkLocation + js_local = js_row_group % max_rg + tid = int(grid.tile_id(np.array(tx), np.array(ty))) + assert tid == js_row_group + assert grid.chunk_location(tid, max_rg) == (js_file, js_local) From cbb7fcdc3c61a6575dafa485412d50c17c7ffd6b Mon Sep 17 00:00:00 2001 From: Nicolas Fernandez Date: Sun, 6 Sep 2026 13:59:20 -0400 Subject: [PATCH 02/17] feat(experimental): add feature catalog and regular-grid Points writer FeatureCatalog assigns genes codes [0, n_genes) in the table's var_names order, so a gene's feature_code is also its CBG row-group index -- no second lookup and no browser-side string join. Non-gene features (controls, unassigned codewords) are retained but coded above every gene and flagged, never folded into a real gene. Verified on Xenium pancreas: 377 genes + 164 controls. write_points_regular_grid keeps every canonical column and the index, and adds display_xy (fixed_size_list[2], already interleaved for deck.gl) and feature_code. Rows are grouped by tile, one row group per logical tile including empty ones, split across zero-padded chunk files. Notable details: - statistics are disabled: the tile formula is the spatial index, so no client reads column-chunk min/max, and they inflate the footer the browser fetches. - coordinates are cast to float64 before the affine: under numpy 2's NEP-50 promotion, float32 * python float stays float32 and loses pixel precision. - the rewrite is idempotent. Re-running on an already-optimized element replaces the render columns rather than appending duplicates, which previously produced an unreadable file ("Multiple matches for FieldRef") because pandas also round-trips fixed_size_list back as a variable-length list. Co-Authored-By: Claude Opus 5 (1M context) --- .../experimental/feature_catalog.py | 194 +++++++++++ .../experimental/points_parquet.py | 289 ++++++++++++++++ tests/test_points_parquet.py | 311 ++++++++++++++++++ 3 files changed, 794 insertions(+) create mode 100644 src/spatialdata_io/experimental/feature_catalog.py create mode 100644 src/spatialdata_io/experimental/points_parquet.py create mode 100644 tests/test_points_parquet.py diff --git a/src/spatialdata_io/experimental/feature_catalog.py b/src/spatialdata_io/experimental/feature_catalog.py new file mode 100644 index 00000000..09ed7d29 --- /dev/null +++ b/src/spatialdata_io/experimental/feature_catalog.py @@ -0,0 +1,194 @@ +"""Stable integer codes for transcript features. + +Rendering a transcript layer needs a compact integer per point, not a string. This module +builds the mapping ``feature name <-> feature_code`` and pins two properties the rest of +the profile depends on: + +1. Genes come first, in the annotating table's ``var_names`` order, so a gene's + ``feature_code`` *is* its row-group index in the cell-by-gene file. No second lookup + table, and no browser-side string join. +2. Non-gene features (negative controls, unassigned codewords) are kept, but are coded + *above* every gene and flagged. They are never silently folded into a real gene, which + would fabricate expression. + +The split point is recorded as ``n_genes`` so a client can tell the two apart from the +manifest alone. +""" + +from __future__ import annotations + +from collections.abc import Iterable, Sequence +from dataclasses import dataclass +from typing import Any + +import numpy as np +import pandas as pd +from numpy.typing import NDArray + +__all__ = ["FeatureCatalog", "CONTROL_PREFIXES"] + +#: Feature-name prefixes that 10x uses for non-gene codewords. Used only for reporting; +#: catalog membership is decided by absence from the table's ``var_names``, not by prefix. +CONTROL_PREFIXES = ( + "NegControlProbe", + "NegControlCodeword", + "UnassignedCodeword", + "antisense", + "BLANK", + "DeprecatedCodeword", + "Intergenic", +) + + +@dataclass(frozen=True) +class FeatureCatalog: + """An ordered feature vocabulary with genes first. + + Parameters + ---------- + names + All feature names. Position in this list is the feature code. + n_genes + Number of leading entries that are genes present in the annotating table. Codes + ``>= n_genes`` are non-gene features. + """ + + names: tuple[str, ...] + n_genes: int + + def __post_init__(self) -> None: + if len(set(self.names)) != len(self.names): + raise ValueError("feature names must be unique") + if not 0 <= self.n_genes <= len(self.names): + raise ValueError(f"n_genes={self.n_genes} out of range for {len(self.names)} features") + + def __len__(self) -> int: + return len(self.names) + + @property + def genes(self) -> tuple[str, ...]: + """The gene names, in table ``var_names`` order.""" + return self.names[: self.n_genes] + + @property + def controls(self) -> tuple[str, ...]: + """The non-gene feature names.""" + return self.names[self.n_genes :] + + @property + def dtype(self) -> np.dtype[Any]: + """Smallest unsigned integer dtype that can hold every code.""" + return np.dtype(np.uint16) if len(self.names) <= np.iinfo(np.uint16).max else np.dtype(np.uint32) + + # -- construction --------------------------------------------------------- + + @classmethod + def from_features_and_table( + cls, + feature_names: Iterable[str], + var_names: Sequence[str], + ) -> FeatureCatalog: + """Build a catalog from the observed transcript features and the table's genes. + + Genes are taken in ``var_names`` order. Any observed feature absent from + ``var_names`` is appended as a control, sorted so the catalog is reproducible. + + Raises + ------ + ValueError + If ``var_names`` contains duplicates, which would make codes ambiguous. + """ + genes = list(var_names) + if len(set(genes)) != len(genes): + raise ValueError("var_names contains duplicates; feature codes would be ambiguous") + + observed = set(feature_names) + gene_set = set(genes) + controls = sorted(observed - gene_set) + return cls(names=tuple(genes) + tuple(controls), n_genes=len(genes)) + + @classmethod + def from_points_and_table( + cls, + points: Any, + table: Any, + feature_key: str = "feature_name", + ) -> FeatureCatalog: + """Build a catalog directly from a Points element and its annotating table. + + Handles the dask categorical read back from a zarr store, whose categories are + lazily "unknown" and raise on access until realized. + """ + col = points[feature_key] + if hasattr(col, "cat"): + try: + features = list(col.cat.categories) + except Exception: + # Unknown categories (typical straight after read_zarr): realize just the + # category list, which is tiny, rather than computing the whole column. + features = list(col.cat.as_known().cat.categories) + else: + features = list(col.unique().compute() if hasattr(col, "compute") else col.unique()) + return cls.from_features_and_table(features, list(table.var_names)) + + # -- encoding ------------------------------------------------------------- + + def encode(self, values: pd.Series | pd.Categorical | NDArray[Any]) -> NDArray[Any]: + """Map feature names to codes. + + Uses the categorical fast path when available, avoiding a per-row dict lookup over + millions of transcripts. + + Raises + ------ + ValueError + If any value is absent from the catalog. Unknown features are never mapped to + a fallback code, since that would silently attribute reads to the wrong feature. + """ + index = pd.Index(self.names) + cat = values.values if isinstance(values, pd.Series) else values + + if isinstance(cat, pd.Categorical): + # Re-map the (small) category list, then take through the codes. + mapped = index.get_indexer(pd.Index(cat.categories)) + self._raise_on_unknown(np.asarray(cat.categories)[mapped < 0]) + codes = np.asarray(cat.codes) + if (codes < 0).any(): + raise ValueError("feature column contains missing (NaN) values") + out = mapped[codes] + else: + arr = np.asarray(cat) + out = index.get_indexer(pd.Index(arr)) + self._raise_on_unknown(np.unique(arr[out < 0])) + + return out.astype(self.dtype, copy=False) + + def _raise_on_unknown(self, unknown: NDArray[Any]) -> None: + if len(unknown): + shown = ", ".join(map(str, unknown[:5])) + more = f" (and {len(unknown) - 5} more)" if len(unknown) > 5 else "" + raise ValueError( + f"{len(unknown)} feature name(s) are not in the catalog: {shown}{more}. " + f"Rebuild the catalog from the same data, or pass the full feature list." + ) + + # -- serialization -------------------------------------------------------- + + def to_frame(self) -> pd.DataFrame: + """Return the catalog as the ``meta_gene.parquet`` table.""" + return pd.DataFrame( + { + "name": list(self.names), + "feature_code": np.arange(len(self.names), dtype=self.dtype), + "is_gene": np.arange(len(self.names)) < self.n_genes, + } + ) + + def to_manifest_dict(self) -> dict[str, Any]: + """Summary for the profile manifest. The full mapping lives in ``meta_gene.parquet``.""" + return { + "n_features": len(self.names), + "n_genes": self.n_genes, + "feature_code_dtype": self.dtype.name, + "gene_codes_match_cbg_row_groups": True, + } diff --git a/src/spatialdata_io/experimental/points_parquet.py b/src/spatialdata_io/experimental/points_parquet.py new file mode 100644 index 00000000..15e34ff5 --- /dev/null +++ b/src/spatialdata_io/experimental/points_parquet.py @@ -0,0 +1,289 @@ +"""Rewrite a Points element into regular-grid row groups. + +The output keeps every canonical column and adds two render-oriented ones: + +``display_xy`` + ``fixed_size_list[2]`` of integer level-0 pixel coordinates. The Arrow child + buffer is therefore already ``[x0, y0, x1, y1, ...]`` -- directly usable as a deck.gl + binary ``getPosition`` attribute with no interleaving step in the browser. +``feature_code`` + Small unsigned integer into the :class:`FeatureCatalog`. + +Physical row order changes (rows are grouped by tile), but no row is added, dropped or +altered, and the DataFrame index is preserved so the reordering is fully traceable. + +Row groups are written one-per-logical-tile *including empty tiles*, so a client can find +a tile's data from the tile formula alone, with no lookup table and no reliance on Parquet +statistics. Output is split across several files because a reader must fetch a whole +footer before reading any row group, and footer size grows with row-group count. +""" + +from __future__ import annotations + +import json +import shutil +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import numpy as np +import pandas as pd +import pyarrow as pa +import pyarrow.parquet as pq +from numpy.typing import NDArray + +from spatialdata_io.experimental.feature_catalog import FeatureCatalog +from spatialdata_io.experimental.regular_grid import ( + DEFAULT_MAX_ROW_GROUPS_PER_FILE, + RegularGrid, +) + +__all__ = [ + "POSITION_COLUMN", + "FEATURE_COLUMN", + "DisplayTransform", + "write_points_regular_grid", +] + +#: Column holding interleaved integer pixel positions. +POSITION_COLUMN = "display_xy" +#: Column holding the integer feature code. +FEATURE_COLUMN = "feature_code" + +#: Largest representable display coordinate. +_UINT32_MAX = np.iinfo(np.uint32).max + + +@dataclass(frozen=True) +class DisplayTransform: + """The affine mapping from canonical element coordinates to display pixels. + + Recorded in the manifest so a client can reproduce the mapping, and so the profile can + be invalidated if the underlying transform changes. + """ + + matrix: tuple[tuple[float, float, float], tuple[float, float, float]] + coordinate_system: str + rounding: str = "nearest" + + @classmethod + def from_element(cls, element: Any, coordinate_system: str) -> DisplayTransform: + """Derive the transform from an element's SpatialData coordinate transformations.""" + from spatialdata.transformations import get_transformation + + t = get_transformation(element, coordinate_system) + affine = np.asarray(t.to_affine_matrix(input_axes=("x", "y"), output_axes=("x", "y"))) + return cls( + matrix=((affine[0, 0], affine[0, 1], affine[0, 2]), (affine[1, 0], affine[1, 1], affine[1, 2])), + coordinate_system=coordinate_system, + ) + + def apply(self, x: NDArray[Any], y: NDArray[Any]) -> tuple[NDArray[np.float64], NDArray[np.float64]]: + """Map canonical coordinates to (unrounded) display pixel coordinates.""" + (a, b, c), (d, e, f) = self.matrix + x = np.asarray(x, dtype=np.float64) + y = np.asarray(y, dtype=np.float64) + return a * x + b * y + c, d * x + e * y + f + + def to_manifest_dict(self) -> dict[str, Any]: + return { + "coordinate_space": "image-pixel", + "coordinate_system": self.coordinate_system, + "affine_matrix": [list(self.matrix[0]), list(self.matrix[1])], + "rounding": self.rounding, + } + + +def _to_display_pixels( + x: NDArray[Any], y: NDArray[Any], transform: DisplayTransform +) -> tuple[NDArray[np.uint32], NDArray[np.uint32]]: + """Transform and round to non-negative integer pixels, validating the declared dtype.""" + px, py = transform.apply(x, y) + + for name, v in (("x", px), ("y", py)): + if not np.isfinite(v).all(): + raise ValueError(f"display {name} contains non-finite values after transform") + + rx = np.rint(px) + ry = np.rint(py) + + for name, v in (("x", rx), ("y", ry)): + lo, hi = float(v.min()), float(v.max()) + if lo < 0: + raise ValueError( + f"display {name} has negative values (min {lo}). display_xy is unsigned; " + f"shift the grid origin or fix the coordinate transform." + ) + if hi > _UINT32_MAX: + raise ValueError(f"display {name} max {hi} exceeds uint32 range") + + return rx.astype(np.uint32), ry.astype(np.uint32) + + +def _interleaved_positions(px: NDArray[np.uint32], py: NDArray[np.uint32]) -> pa.FixedSizeListArray: + """Build ``fixed_size_list[2]`` whose child buffer is ``[x0,y0,x1,y1,...]``.""" + flat = np.empty(px.size * 2, dtype=np.uint32) + flat[0::2] = px + flat[1::2] = py + return pa.FixedSizeListArray.from_arrays(pa.array(flat), 2) + + +def write_points_regular_grid( + points: Any, + output_dir: str | Path, + *, + catalog: FeatureCatalog, + grid: RegularGrid | None = None, + coordinate_system: str = "global", + feature_key: str = "feature_name", + tile_size_px: float = 250.0, + max_row_groups_per_file: int = DEFAULT_MAX_ROW_GROUPS_PER_FILE, + compression: str = "snappy", + overwrite: bool = False, +) -> dict[str, Any]: + """Write a Points element as regular-grid row groups. + + Parameters + ---------- + points + A SpatialData Points element (dask DataFrame) or a pandas DataFrame. + output_dir + Directory to write the chunk files into. Written atomically: a temporary sibling + directory is populated first and swapped in only on success. + catalog + Feature vocabulary providing ``feature_code``. + grid + The tile grid. If ``None``, the smallest grid covering the data is derived using + ``tile_size_px``. + coordinate_system + SpatialData coordinate system defining display pixel space. + feature_key + Column holding the feature name. + tile_size_px + Tile edge length, used only when ``grid`` is ``None``. + max_row_groups_per_file + Row groups per chunk file. + compression + Parquet compression codec. + overwrite + Replace ``output_dir`` if it exists. + + Returns + ------- + The manifest fragment describing the written files. + """ + output_dir = Path(output_dir) + if output_dir.exists() and not overwrite: + raise FileExistsError(f"{output_dir} exists; pass overwrite=True to replace it") + + df = points.compute() if hasattr(points, "compute") else points + if not isinstance(df, pd.DataFrame): + raise TypeError(f"expected a DataFrame, got {type(df).__name__}") + if feature_key not in df.columns: + raise ValueError(f"feature column {feature_key!r} not found; have {list(df.columns)}") + for axis in ("x", "y"): + if axis not in df.columns: + raise ValueError(f"points element has no {axis!r} column; have {list(df.columns)}") + + transform = DisplayTransform.from_element(points, coordinate_system) + px, py = _to_display_pixels(df["x"].to_numpy(), df["y"].to_numpy(), transform) + + if grid is None: + grid = RegularGrid.from_bounds(0, 0, float(px.max()), float(py.max()), tile_size_px) + + tile_ids = grid.assign(px, py) + codes = catalog.encode(df[feature_key]) + + # Keep every canonical column and index; append the two render columns. + # The transform lives in .attrs and is not JSON-serializable, so drop it before the + # Arrow conversion exactly as spatialdata's own points writer does -- it is persisted + # in the element's zarr attributes, not in the parquet file. + stale = [c for c in (POSITION_COLUMN, FEATURE_COLUMN) if c in df.columns] + if df.attrs or stale: + df = df.copy(deep=False) + df.attrs = {} + # Re-running the optimizer on an already-optimized element must replace the render + # columns, not append duplicates. A duplicated name makes the file unreadable by + # column projection ("Multiple matches for FieldRef"), and pandas round-trips + # fixed_size_list back as a variable-length list, so the stale copy is also the + # wrong Arrow type. Both are recomputed below from the canonical coordinates. + if stale: + df = df.drop(columns=stale) + table = pa.Table.from_pandas(df, preserve_index=True) + table = table.append_column(POSITION_COLUMN, _interleaved_positions(px, py)) + table = table.append_column(FEATURE_COLUMN, pa.array(codes)) + + # Stable sort keeps the original relative order inside a tile, so the rewrite is + # deterministic and diffable. + order = np.argsort(tile_ids, kind="stable") + table = table.take(pa.array(order)) + sorted_tile_ids = tile_ids[order] + + # Row-group boundaries: offsets[t]..offsets[t+1] is tile t's slice. + offsets = np.searchsorted(sorted_tile_ids, np.arange(grid.num_tiles + 1), side="left") + + staging = output_dir.with_name(output_dir.name + ".tmp") + if staging.exists(): + shutil.rmtree(staging) + staging.mkdir(parents=True) + + filenames = grid.chunk_filenames(max_row_groups_per_file) + schema = table.schema.with_metadata( + { + **(table.schema.metadata or {}), + b"profile": b"celldega_regular_grid_v1", + b"storage_mode": b"row_groups_chunked", + b"max_row_groups_per_file": str(max_row_groups_per_file).encode(), + b"tile_grid": json.dumps(grid.to_manifest_dict()).encode(), + } + ) + + try: + writer: pq.ParquetWriter | None = None + current_file = -1 + for tile_id in range(grid.num_tiles): + file_index, _ = grid.chunk_location(tile_id, max_row_groups_per_file) + if file_index != current_file: + if writer is not None: + writer.close() + writer = pq.ParquetWriter( + staging / filenames[file_index], + schema, + compression=compression, + # Statistics are dead weight here: the tile formula is the spatial + # index, so no client ever consults per-column-chunk min/max, and + # they inflate the footer the browser must download up front. + write_statistics=False, + ) + current_file = file_index + + start, end = int(offsets[tile_id]), int(offsets[tile_id + 1]) + assert writer is not None + # Empty tiles are written as zero-row row groups so that + # row_group_index == tile_id holds without a lookup table. + writer.write_table(table.slice(start, end - start) if end > start else schema.empty_table()) + if writer is not None: + writer.close() + except BaseException: + shutil.rmtree(staging, ignore_errors=True) + raise + + if output_dir.exists(): + shutil.rmtree(output_dir) + staging.rename(output_dir) + + return { + "directory": str(output_dir.name), + "files": filenames, + "max_row_groups_per_file": max_row_groups_per_file, + "total_row_groups": grid.num_tiles, + "position_column": POSITION_COLUMN, + "position_encoding": "fixed_size_list", + "position_dtype": "uint32", + "position_size": 2, + "feature_column": FEATURE_COLUMN, + "n_rows": int(table.num_rows), + "tile_grid": grid.to_manifest_dict(), + "display_transform": transform.to_manifest_dict(), + "feature_catalog": catalog.to_manifest_dict(), + } diff --git a/tests/test_points_parquet.py b/tests/test_points_parquet.py new file mode 100644 index 00000000..18d93dd2 --- /dev/null +++ b/tests/test_points_parquet.py @@ -0,0 +1,311 @@ +"""Tests for the regular-grid Points rewrite. + +The synthetic fixture is a 2x3 tile grid deliberately containing the awkward cases: +an empty tile, a point exactly on a tile boundary, several points sharing one tile, +and both genes and a non-gene control feature. +""" + +from __future__ import annotations + +from pathlib import Path + +import numpy as np +import pandas as pd +import pyarrow as pa +import pyarrow.parquet as pq +import pytest + +from spatialdata_io.experimental.feature_catalog import FeatureCatalog +from spatialdata_io.experimental.points_parquet import ( + FEATURE_COLUMN, + POSITION_COLUMN, + write_points_regular_grid, +) +from spatialdata_io.experimental.regular_grid import RegularGrid + +# Display pixel space is 20 x 30 px => a 2 x 3 grid of 10 px tiles. +# Canonical (micron) coordinates are half of the pixel values, i.e. Scale([2, 2]). +TILE_SIZE = 10.0 +GRID = RegularGrid(origin_x=0.0, origin_y=0.0, tile_size_px=TILE_SIZE, num_tiles_x=2, num_tiles_y=3) + +#: (pixel_x, pixel_y, feature, expected_tile_id). Tile 4 is intentionally absent. +POINTS_SPEC = [ + (1.0, 1.0, "GENEA", 0), + (2.0, 3.0, "GENEB", 0), + (5.0, 5.0, "GENEA", 0), # three points share tile 0 + (3.0, 12.0, "GENEB", 1), + (4.0, 25.0, "GENEA", 2), + (10.0, 2.0, "GENEB", 3), # exactly on the x tile boundary -> upper tile + (15.0, 28.0, "NegControlProbe_00042", 5), +] +VAR_NAMES = ["GENEA", "GENEB"] + + +@pytest.fixture +def catalog() -> FeatureCatalog: + return FeatureCatalog.from_features_and_table([s[2] for s in POINTS_SPEC], VAR_NAMES) + + +@pytest.fixture +def points() -> pd.DataFrame: + """A pandas Points-like frame with canonical micron coords and extra annotations.""" + px = np.array([s[0] for s in POINTS_SPEC]) + py = np.array([s[1] for s in POINTS_SPEC]) + df = pd.DataFrame( + { + "x": px / 2.0, + "y": py / 2.0, + "z": np.linspace(0.0, 1.0, len(POINTS_SPEC)), + "feature_name": pd.Categorical([s[2] for s in POINTS_SPEC]), + "cell_id": [f"cell-{i}" for i in range(len(POINTS_SPEC))], + "transcript_id": np.arange(100, 100 + len(POINTS_SPEC), dtype=np.uint64), + "qv": np.linspace(20.0, 40.0, len(POINTS_SPEC)).astype(np.float32), + } + ) + from spatialdata.models import PointsModel + + return PointsModel.parse(df, coordinates={"x": "x", "y": "y", "z": "z"}, feature_key="feature_name") + + +@pytest.fixture +def written(tmp_path: Path, points: pd.DataFrame, catalog: FeatureCatalog) -> tuple[Path, dict]: + from spatialdata.transformations import Scale, set_transformation + + set_transformation(points, Scale([2.0, 2.0], axes=("x", "y")), "global") + out = tmp_path / "points.parquet" + manifest = write_points_regular_grid(points, out, catalog=catalog, grid=GRID) + return out, manifest + + +def _read_all(directory: Path, manifest: dict) -> pa.Table: + return pa.concat_tables([pq.read_table(directory / f) for f in manifest["files"]]) + + +# -- row-group layout --------------------------------------------------------- + + +def test_row_group_count_equals_tile_count(written: tuple[Path, dict]) -> None: + directory, manifest = written + total = sum(pq.ParquetFile(directory / f).metadata.num_row_groups for f in manifest["files"]) + assert total == GRID.num_tiles == 6 + assert manifest["total_row_groups"] == 6 + + +def test_each_tile_row_group_holds_exactly_its_points(written: tuple[Path, dict]) -> None: + """The core contract: row_group_index == tile_id, with no lookup table.""" + directory, manifest = written + expected: dict[int, list[tuple[float, float]]] = {} + for pxv, pyv, _, tid in POINTS_SPEC: + expected.setdefault(tid, []).append((pxv, pyv)) + + for tile_id in range(GRID.num_tiles): + file_index, local = GRID.chunk_location(tile_id, manifest["max_row_groups_per_file"]) + f = pq.ParquetFile(directory / manifest["files"][file_index]) + rg = f.read_row_group(local, columns=[POSITION_COLUMN]) + got = [tuple(v) for v in rg[POSITION_COLUMN].to_pylist()] + assert sorted(got) == sorted(expected.get(tile_id, [])), f"tile {tile_id}" + + +def test_empty_tile_is_a_zero_row_row_group(written: tuple[Path, dict]) -> None: + directory, manifest = written + file_index, local = GRID.chunk_location(4, manifest["max_row_groups_per_file"]) + f = pq.ParquetFile(directory / manifest["files"][file_index]) + assert f.metadata.row_group(local).num_rows == 0 + assert f.read_row_group(local).num_rows == 0 + + +def test_boundary_point_goes_to_the_upper_tile(written: tuple[Path, dict]) -> None: + """A point at exactly x=10 belongs to tile_x=1 under half-open bounds.""" + directory, manifest = written + f = pq.ParquetFile(directory / manifest["files"][0]) + rg = f.read_row_group(3, columns=[POSITION_COLUMN]) + assert [tuple(v) for v in rg[POSITION_COLUMN].to_pylist()] == [(10, 2)] + + +# -- canonical data preservation ---------------------------------------------- + + +def test_no_row_lost_or_duplicated(written: tuple[Path, dict], points: pd.DataFrame) -> None: + directory, manifest = written + table = _read_all(directory, manifest) + original = points.compute() if hasattr(points, "compute") else points + assert table.num_rows == len(POINTS_SPEC) + assert sorted(table["transcript_id"].to_pylist()) == sorted(original["transcript_id"].tolist()) + + +def test_canonical_columns_are_unchanged(written: tuple[Path, dict], points: pd.DataFrame) -> None: + """Values must survive the reorder exactly; only row order may differ.""" + directory, manifest = written + got = _read_all(directory, manifest).to_pandas() + original = points.compute() if hasattr(points, "compute") else points + + merged = got.set_index("transcript_id").loc[original["transcript_id"].to_numpy()] + for col in ("x", "y", "z", "cell_id", "qv"): + np.testing.assert_array_equal( + merged[col].to_numpy(), original[col].to_numpy(), err_msg=f"column {col} changed" + ) + assert list(merged["feature_name"].astype(str)) == list(original["feature_name"].astype(str)) + + +def test_index_is_preserved(written: tuple[Path, dict], points: pd.DataFrame) -> None: + directory, manifest = written + got = _read_all(directory, manifest).to_pandas() + original = points.compute() if hasattr(points, "compute") else points + assert sorted(got.index.tolist()) == sorted(original.index.tolist()) + + +def test_row_order_is_grouped_by_tile(written: tuple[Path, dict]) -> None: + directory, manifest = written + table = _read_all(directory, manifest) + xs = np.array([v[0] for v in table[POSITION_COLUMN].to_pylist()]) + ys = np.array([v[1] for v in table[POSITION_COLUMN].to_pylist()]) + tile_ids = GRID.assign(xs, ys) + assert (np.diff(tile_ids) >= 0).all(), "rows are not grouped by tile" + + +# -- render columns ----------------------------------------------------------- + + +def test_display_xy_is_fixed_size_list_uint32(written: tuple[Path, dict]) -> None: + directory, manifest = written + field = pq.ParquetFile(directory / manifest["files"][0]).schema_arrow.field(POSITION_COLUMN) + assert field.type == pa.list_(pa.uint32(), 2) + assert manifest["position_dtype"] == "uint32" + assert manifest["position_encoding"] == "fixed_size_list" + + +def test_display_xy_child_buffer_is_interleaved(written: tuple[Path, dict]) -> None: + """The flat child buffer must be [x0,y0,x1,y1,...] so deck.gl can consume it directly.""" + directory, manifest = written + col = _read_all(directory, manifest)[POSITION_COLUMN].combine_chunks() + flat = col.values.to_numpy(zero_copy_only=False) + pairs = [tuple(v) for v in col.to_pylist()] + assert flat[0::2].tolist() == [p[0] for p in pairs] + assert flat[1::2].tolist() == [p[1] for p in pairs] + + +def test_display_xy_matches_the_transform(written: tuple[Path, dict]) -> None: + directory, manifest = written + table = _read_all(directory, manifest) + got = {int(t): tuple(v) for t, v in zip(table["transcript_id"].to_pylist(), table[POSITION_COLUMN].to_pylist())} + for i, (pxv, pyv, _, _) in enumerate(POINTS_SPEC): + assert got[100 + i] == (int(pxv), int(pyv)) + + +def test_feature_codes_match_catalog(written: tuple[Path, dict], catalog: FeatureCatalog) -> None: + directory, manifest = written + table = _read_all(directory, manifest) + for name, code in zip(table["feature_name"].to_pylist(), table[FEATURE_COLUMN].to_pylist()): + assert catalog.names[code] == name + + +def test_control_feature_is_coded_above_every_gene(written: tuple[Path, dict], catalog: FeatureCatalog) -> None: + directory, manifest = written + table = _read_all(directory, manifest).to_pandas() + control = table[table["feature_name"].astype(str).str.startswith("NegControl")] + assert len(control) == 1 + assert int(control[FEATURE_COLUMN].iloc[0]) >= catalog.n_genes + + +# -- column projection -------------------------------------------------------- + + +def test_column_projection_reads_only_render_columns(written: tuple[Path, dict]) -> None: + """Celldega projects these two columns; canonical columns must not be required.""" + directory, manifest = written + f = pq.ParquetFile(directory / manifest["files"][0]) + rg = f.read_row_group(0, columns=[POSITION_COLUMN, FEATURE_COLUMN]) + assert rg.column_names == [POSITION_COLUMN, FEATURE_COLUMN] + assert rg.num_rows == 3 + + +# -- file layout -------------------------------------------------------------- + + +def test_multi_file_split_and_padding(tmp_path: Path, points: pd.DataFrame, catalog: FeatureCatalog) -> None: + from spatialdata.transformations import Scale, set_transformation + + set_transformation(points, Scale([2.0, 2.0], axes=("x", "y")), "global") + out = tmp_path / "multi.parquet" + manifest = write_points_regular_grid(points, out, catalog=catalog, grid=GRID, max_row_groups_per_file=2) + + assert manifest["files"] == ["chunk_0.parquet", "chunk_1.parquet", "chunk_2.parquet"] + assert [pq.ParquetFile(out / f).metadata.num_row_groups for f in manifest["files"]] == [2, 2, 2] + assert sum(pq.ParquetFile(out / f).metadata.num_rows for f in manifest["files"]) == len(POINTS_SPEC) + + +def test_statistics_are_disabled(written: tuple[Path, dict]) -> None: + """Footer weight matters in the browser and the tile formula is the spatial index.""" + directory, manifest = written + rg = pq.ParquetFile(directory / manifest["files"][0]).metadata.row_group(0) + assert not rg.column(0).is_stats_set + + +def test_overwrite_guard(written: tuple[Path, dict], points: pd.DataFrame, catalog: FeatureCatalog) -> None: + directory, _ = written + with pytest.raises(FileExistsError): + write_points_regular_grid(points, directory, catalog=catalog, grid=GRID) + + +def test_failure_leaves_no_partial_output(tmp_path: Path, points: pd.DataFrame) -> None: + """A mid-write error must not leave a half-rewritten directory behind.""" + from spatialdata.transformations import Scale, set_transformation + + set_transformation(points, Scale([2.0, 2.0], axes=("x", "y")), "global") + out = tmp_path / "boom.parquet" + bad = FeatureCatalog(names=("GENEA",), n_genes=1) # missing GENEB -> encode() raises + with pytest.raises(ValueError, match="not in the catalog"): + write_points_regular_grid(points, out, catalog=bad, grid=GRID) + assert not out.exists() + assert not out.with_name(out.name + ".tmp").exists() + + +# -- validation --------------------------------------------------------------- + + +def test_negative_display_coordinates_are_rejected(tmp_path: Path, catalog: FeatureCatalog) -> None: + from spatialdata.models import PointsModel + from spatialdata.transformations import Scale, set_transformation + + df = pd.DataFrame({"x": [-5.0, 1.0], "y": [1.0, 1.0], "feature_name": pd.Categorical(["GENEA", "GENEB"])}) + p = PointsModel.parse(df, coordinates={"x": "x", "y": "y"}, feature_key="feature_name") + set_transformation(p, Scale([2.0, 2.0], axes=("x", "y")), "global") + with pytest.raises(ValueError, match="negative values"): + write_points_regular_grid(p, tmp_path / "neg.parquet", catalog=catalog, grid=GRID) + + +def test_missing_feature_column_is_reported(tmp_path: Path, points: pd.DataFrame, catalog: FeatureCatalog) -> None: + with pytest.raises(ValueError, match="feature column 'nope' not found"): + write_points_regular_grid(points, tmp_path / "x.parquet", catalog=catalog, grid=GRID, feature_key="nope") + + +def test_rewrite_is_idempotent(tmp_path: Path, points: pd.DataFrame, catalog: FeatureCatalog) -> None: + """Re-optimizing an already-optimized element replaces the render columns, not appends them. + + A duplicated column name makes projected reads fail outright, and pandas round-trips + fixed_size_list back as a variable-length list, so a stale copy is also mistyped. + """ + from spatialdata.transformations import Scale, set_transformation + + set_transformation(points, Scale([2.0, 2.0], axes=("x", "y")), "global") + out = tmp_path / "idem.parquet" + m1 = write_points_regular_grid(points, out, catalog=catalog, grid=GRID) + first = _read_all(out, m1) + + # Feed the written result back in, the way read_zarr would hand it back: a dask + # frame that already carries display_xy/feature_code. + import dask.dataframe as dd + + again = dd.from_pandas(first.to_pandas(), npartitions=1) + # Attach the transform the way read_zarr does, rather than via set_transformation, + # which requires an element that already carries one. + again.attrs["transform"] = {"global": Scale([2.0, 2.0], axes=("x", "y"))} + assert POSITION_COLUMN in again.columns # precondition: the stale columns are present + m2 = write_points_regular_grid(again, out, catalog=catalog, grid=GRID, overwrite=True) + second = _read_all(out, m2) + + assert second.column_names.count(POSITION_COLUMN) == 1 + assert second.column_names.count(FEATURE_COLUMN) == 1 + assert second.schema.field(POSITION_COLUMN).type == pa.list_(pa.uint32(), 2) + assert second.num_rows == first.num_rows + assert second[POSITION_COLUMN].to_pylist() == first[POSITION_COLUMN].to_pylist() From 3b2b45ffb8a86118e3865166390e2b35d22b42d3 Mon Sep 17 00:00:00 2001 From: Nicolas Fernandez Date: Sun, 6 Sep 2026 14:16:08 -0400 Subject: [PATCH 03/17] feat(experimental): add regular-grid Shapes writer Adds display_geometry (list[2]>>: polygon -> rings -> interleaved integer pixel vertices) and cell_code, alongside the untouched canonical WKB geometry. - Cells are assigned to exactly one tile by centroid in display pixel space; a polygon whose outline crosses a tile boundary is not duplicated. Covered by a fixture case with exactly that shape. - display_geometry is marked lossy in the manifest: exterior ring only, largest part of a MultiPolygon, matching Celldega's current behaviour. - The nested layout is walked by a test that mirrors getPolygonDataFromChunk (polygon offset -> ring offset -> coordinate index), so a layout change that would break the browser fails here instead. - GeoParquet 'geo' metadata is preserved via the same conversion to_parquet uses. GeoDataFrame.to_arrow() emits GeoArrow extension metadata but not the 'geo' key, and without it the rewritten file stops being readable by geopandas.read_parquet and by SpatialData. Verified against Xenium pancreas: 140,702 cells, all single-ring Polygons, shapes index matches table obs_names. Co-Authored-By: Claude Opus 5 (1M context) --- .../experimental/points_parquet.py | 5 +- .../experimental/shapes_parquet.py | 252 ++++++++++++++++++ tests/test_shapes_parquet.py | 219 +++++++++++++++ 3 files changed, 475 insertions(+), 1 deletion(-) create mode 100644 src/spatialdata_io/experimental/shapes_parquet.py create mode 100644 tests/test_shapes_parquet.py diff --git a/src/spatialdata_io/experimental/points_parquet.py b/src/spatialdata_io/experimental/points_parquet.py index 15e34ff5..65c30e90 100644 --- a/src/spatialdata_io/experimental/points_parquet.py +++ b/src/spatialdata_io/experimental/points_parquet.py @@ -134,6 +134,7 @@ def write_points_regular_grid( *, catalog: FeatureCatalog, grid: RegularGrid | None = None, + display_transform: DisplayTransform | None = None, coordinate_system: str = "global", feature_key: str = "feature_name", tile_size_px: float = 250.0, @@ -185,7 +186,9 @@ def write_points_regular_grid( if axis not in df.columns: raise ValueError(f"points element has no {axis!r} column; have {list(df.columns)}") - transform = DisplayTransform.from_element(points, coordinate_system) + # When called as a SpatialData ``points_writer`` hook the element arrives with its + # transformations already stripped from attrs, so the caller must supply the transform. + transform = display_transform or DisplayTransform.from_element(points, coordinate_system) px, py = _to_display_pixels(df["x"].to_numpy(), df["y"].to_numpy(), transform) if grid is None: diff --git a/src/spatialdata_io/experimental/shapes_parquet.py b/src/spatialdata_io/experimental/shapes_parquet.py new file mode 100644 index 00000000..208f6953 --- /dev/null +++ b/src/spatialdata_io/experimental/shapes_parquet.py @@ -0,0 +1,252 @@ +"""Rewrite a Shapes element into regular-grid row groups. + +Adds two render-oriented columns beside the canonical geometry: + +``display_geometry`` + ``list[2]>>`` -- polygon -> rings -> interleaved integer + pixel vertices. The nesting is chosen so a client can lift deck.gl's ``getPolygon`` + straight out of the flat coordinate child buffer and ``startIndices`` out of the list + offsets, with no WKB parsing and no per-vertex JavaScript objects. +``cell_code`` + Positional index into the annotating table, so cells can be coloured from a + cell-by-gene vector without a string join in the browser. + +``display_geometry`` is explicitly a *lossy display* representation: only the exterior +ring is kept, and for a MultiPolygon only its largest part. The canonical geometry column +is written through unchanged, and the GeoParquet metadata is preserved so the file is +still readable by :func:`geopandas.read_parquet` and by SpatialData itself. + +Each cell is assigned to exactly one tile, by its centroid in display pixel space. A +polygon whose outline crosses into a neighbouring tile is *not* duplicated -- duplicating +would inflate the file and make cell counts wrong. +""" + +from __future__ import annotations + +import json +import shutil +from pathlib import Path +from typing import Any + +import numpy as np +import pyarrow as pa +import pyarrow.parquet as pq +import shapely +from numpy.typing import NDArray + +from spatialdata_io.experimental.points_parquet import DisplayTransform, _to_display_pixels +from spatialdata_io.experimental.regular_grid import ( + DEFAULT_MAX_ROW_GROUPS_PER_FILE, + RegularGrid, +) + +__all__ = ["GEOMETRY_COLUMN", "CELL_CODE_COLUMN", "write_shapes_regular_grid"] + +#: Column holding the nested integer-pixel display polygons. +GEOMETRY_COLUMN = "display_geometry" +#: Column holding the positional cell index. +CELL_CODE_COLUMN = "cell_code" + + +def _largest_polygon(geom: Any) -> Any: + """Reduce a MultiPolygon to its largest part, matching Celldega's existing behaviour.""" + if geom is None or geom.is_empty: + return geom + if geom.geom_type == "MultiPolygon": + return max(geom.geoms, key=lambda g: g.area) + return geom + + +def _exterior_only(geometry: Any) -> NDArray[Any]: + """Return an array of single-ring polygons: largest part, exterior ring only.""" + geoms = np.asarray([_largest_polygon(g) for g in geometry], dtype=object) + rings = shapely.get_exterior_ring(geoms) + return shapely.polygons(rings) + + +def _display_geometry_array( + geometry: Any, transform: DisplayTransform +) -> tuple[pa.ListArray, NDArray[np.uint32], NDArray[np.uint32]]: + """Build the nested display-geometry array and the per-cell display centroids. + + Returns the Arrow array plus the integer pixel centroid coordinates used for tiling. + """ + simple = _exterior_only(geometry) + + # to_ragged_array gives exactly the buffers the nested Arrow layout needs: + # a flat (N, 2) coordinate array plus ring and polygon offsets. + _, coords, offsets = shapely.to_ragged_array(simple) + ring_offsets, polygon_offsets = offsets + + px, py = _to_display_pixels(coords[:, 0], coords[:, 1], transform) + flat = np.empty(px.size * 2, dtype=np.uint32) + flat[0::2] = px + flat[1::2] = py + + vertices = pa.FixedSizeListArray.from_arrays(pa.array(flat), 2) + rings = pa.ListArray.from_arrays(pa.array(ring_offsets, type=pa.int32()), vertices) + polygons = pa.ListArray.from_arrays(pa.array(polygon_offsets, type=pa.int32()), rings) + + centroids = shapely.centroid(simple) + cx, cy = _to_display_pixels(shapely.get_x(centroids), shapely.get_y(centroids), transform) + return polygons, cx, cy + + +def _canonical_geoparquet_table(shapes: Any) -> pa.Table: + """Convert a GeoDataFrame to Arrow while keeping the GeoParquet ``geo`` metadata. + + ``GeoDataFrame.to_arrow()`` emits GeoArrow *extension* metadata but not the GeoParquet + ``geo`` schema key, which is added by ``to_parquet()``. Without that key the rewritten + file is no longer readable by :func:`geopandas.read_parquet` or by SpatialData. We + therefore reuse the same conversion ``to_parquet`` performs, rather than writing the + file once just to recover its metadata. + """ + try: + from geopandas.io.arrow import _geopandas_to_arrow + except ImportError as exc: # pragma: no cover - depends on geopandas internals + raise RuntimeError( + "could not access geopandas' arrow conversion; a geopandas version with " + "geopandas.io.arrow._geopandas_to_arrow is required to preserve GeoParquet metadata" + ) from exc + + table = _geopandas_to_arrow(shapes, index=None, geometry_encoding="WKB") + if b"geo" not in (table.schema.metadata or {}): # pragma: no cover - defensive + raise RuntimeError("geopandas did not produce GeoParquet 'geo' metadata") + return table + + +def write_shapes_regular_grid( + shapes: Any, + output_path: str | Path, + *, + grid: RegularGrid, + display_transform: DisplayTransform | None = None, + coordinate_system: str = "global", + cell_index: Any | None = None, + max_row_groups_per_file: int = DEFAULT_MAX_ROW_GROUPS_PER_FILE, + compression: str = "snappy", + overwrite: bool = False, +) -> dict[str, Any]: + """Write a Shapes element as regular-grid row groups. + + Parameters + ---------- + shapes + A SpatialData Shapes element (GeoDataFrame). + output_path + Destination. A single ``shapes.parquet`` file when it fits in one chunk (which is + what SpatialData's reader expects), otherwise a directory of chunk files. + grid + The tile grid; must be the same grid used for the points element. + display_transform + Transform to display pixel space. Derived from the element when omitted. + coordinate_system + Coordinate system used when deriving the transform. + cell_index + Index defining ``cell_code`` order, normally the annotating table's ``obs_names``. + Defaults to the shapes' own index order. + max_row_groups_per_file + Row groups per chunk file. + compression + Parquet compression codec. + overwrite + Replace an existing output. + + Returns + ------- + The manifest fragment describing the written file(s). + """ + output_path = Path(output_path) + if output_path.exists() and not overwrite: + raise FileExistsError(f"{output_path} exists; pass overwrite=True to replace it") + + transform = display_transform or DisplayTransform.from_element(shapes, coordinate_system) + display, cx, cy = _display_geometry_array(shapes.geometry, transform) + + table = _canonical_geoparquet_table(shapes) + + if cell_index is None: + codes = np.arange(len(shapes), dtype=np.uint32) + else: + positions = {k: i for i, k in enumerate(cell_index)} + missing = [k for k in shapes.index if k not in positions] + if missing: + raise ValueError( + f"{len(missing)} shape(s) are absent from cell_index (e.g. {missing[:3]}); " + f"cell_code would be undefined. Pass the table's obs_names for these cells." + ) + codes = np.fromiter((positions[k] for k in shapes.index), dtype=np.uint32, count=len(shapes)) + + table = table.append_column(GEOMETRY_COLUMN, display) + table = table.append_column(CELL_CODE_COLUMN, pa.array(codes)) + + tile_ids = grid.assign(cx, cy) + order = np.argsort(tile_ids, kind="stable") + table = table.take(pa.array(order)) + offsets = np.searchsorted(tile_ids[order], np.arange(grid.num_tiles + 1), side="left") + + n_files = grid.num_files(max_row_groups_per_file) + single_file = n_files == 1 + filenames = ["shapes.parquet"] if single_file else grid.chunk_filenames(max_row_groups_per_file) + + schema = table.schema.with_metadata( + { + **(table.schema.metadata or {}), + b"profile": b"celldega_regular_grid_v1", + b"storage_mode": b"row_groups_chunked", + b"max_row_groups_per_file": str(max_row_groups_per_file).encode(), + b"tile_grid": json.dumps(grid.to_manifest_dict()).encode(), + } + ) + + staging = output_path.with_name(output_path.name + ".tmp") + if staging.exists(): + shutil.rmtree(staging) if staging.is_dir() else staging.unlink() + + try: + if single_file: + staging.parent.mkdir(parents=True, exist_ok=True) + targets = [staging] + else: + staging.mkdir(parents=True) + targets = [staging / f for f in filenames] + + writer: pq.ParquetWriter | None = None + current = -1 + for tile_id in range(grid.num_tiles): + file_index = 0 if single_file else grid.chunk_location(tile_id, max_row_groups_per_file)[0] + if file_index != current: + if writer is not None: + writer.close() + writer = pq.ParquetWriter(targets[file_index], schema, compression=compression, write_statistics=False) + current = file_index + start, end = int(offsets[tile_id]), int(offsets[tile_id + 1]) + assert writer is not None + writer.write_table(table.slice(start, end - start) if end > start else schema.empty_table()) + if writer is not None: + writer.close() + except BaseException: + if staging.exists(): + shutil.rmtree(staging) if staging.is_dir() else staging.unlink() + raise + + if output_path.exists(): + shutil.rmtree(output_path) if output_path.is_dir() else output_path.unlink() + staging.rename(output_path) + + fragment: dict[str, Any] = { + "geometry_column": GEOMETRY_COLUMN, + "cell_id_column": CELL_CODE_COLUMN, + "max_row_groups_per_file": max_row_groups_per_file, + "total_row_groups": grid.num_tiles, + "n_shapes": int(table.num_rows), + "geometry_is_lossy": True, + "geometry_note": "exterior ring of the largest polygon part; canonical geometry retained", + "tile_grid": grid.to_manifest_dict(), + } + if single_file: + fragment["path"] = output_path.name + else: + fragment["directory"] = output_path.name + fragment["files"] = filenames + return fragment diff --git a/tests/test_shapes_parquet.py b/tests/test_shapes_parquet.py new file mode 100644 index 00000000..16bf68e5 --- /dev/null +++ b/tests/test_shapes_parquet.py @@ -0,0 +1,219 @@ +"""Tests for the regular-grid Shapes rewrite. + +The fixture deliberately includes a polygon whose outline crosses a tile boundary while +its centroid does not, since assigning by centroid (rather than by overlap) is what keeps +each cell in exactly one row group. +""" + +from __future__ import annotations + +from pathlib import Path + +import geopandas as gpd +import numpy as np +import pyarrow as pa +import pyarrow.parquet as pq +import pytest +from shapely.geometry import MultiPolygon, Polygon + +from spatialdata_io.experimental.points_parquet import DisplayTransform +from spatialdata_io.experimental.regular_grid import RegularGrid +from spatialdata_io.experimental.shapes_parquet import ( + CELL_CODE_COLUMN, + GEOMETRY_COLUMN, + write_shapes_regular_grid, +) + +# Same 2x3 grid of 10 px tiles as the points fixture; canonical coords are half of pixels. +GRID = RegularGrid(origin_x=0.0, origin_y=0.0, tile_size_px=10.0, num_tiles_x=2, num_tiles_y=3) +XFORM = DisplayTransform(matrix=((2.0, 0.0, 0.0), (0.0, 2.0, 0.0)), coordinate_system="global") + + +def _square(cx: float, cy: float, half: float) -> Polygon: + """A square in *canonical* coords, centred on (cx, cy).""" + return Polygon( + [(cx - half, cy - half), (cx + half, cy - half), (cx + half, cy + half), (cx - half, cy + half)] + ) + + +#: name -> (geometry, expected tile id). Canonical coords; pixels are 2x. +SHAPES_SPEC = { + "cell-a": (_square(2.5, 2.5, 1.0), 0), # centroid px (5,5) -> tile (0,0) + "cell-b": (_square(2.0, 7.5, 1.0), 1), # centroid px (4,15) -> tile (0,1) + "cell-c": (_square(2.0, 12.5, 1.0), 2), # centroid px (4,25) -> tile (0,2) + # centroid px (8,4) -> tile (0,0), but the outline reaches x=12px, crossing into tile_x=1 + "cell-crossing": (_square(4.0, 2.0, 2.0), 0), + "cell-e": (_square(7.5, 12.5, 1.0), 5), # centroid px (15,25) -> tile (1,2) +} + + +@pytest.fixture +def shapes() -> gpd.GeoDataFrame: + from spatialdata.models import ShapesModel + + gdf = gpd.GeoDataFrame( + {"area": [g.area for g, _ in SHAPES_SPEC.values()]}, + geometry=[g for g, _ in SHAPES_SPEC.values()], + index=list(SHAPES_SPEC), + ) + return ShapesModel.parse(gdf) + + +@pytest.fixture +def written(tmp_path: Path, shapes: gpd.GeoDataFrame) -> tuple[Path, dict]: + out = tmp_path / "shapes.parquet" + manifest = write_shapes_regular_grid(shapes, out, grid=GRID, display_transform=XFORM) + return out, manifest + + +# -- layout ------------------------------------------------------------------- + + +def test_single_file_when_grid_fits_one_chunk(written: tuple[Path, dict]) -> None: + out, manifest = written + assert out.is_file() + assert manifest["path"] == "shapes.parquet" + assert pq.ParquetFile(out).metadata.num_row_groups == GRID.num_tiles + + +def test_each_cell_is_in_its_centroid_tile(written: tuple[Path, dict]) -> None: + out, _ = written + f = pq.ParquetFile(out) + for tile_id in range(GRID.num_tiles): + expected = sorted(n for n, (_, t) in SHAPES_SPEC.items() if t == tile_id) + rg = f.read_row_group(tile_id, columns=[CELL_CODE_COLUMN]) + got = sorted(list(SHAPES_SPEC)[c] for c in rg[CELL_CODE_COLUMN].to_pylist()) + assert got == expected, f"tile {tile_id}" + + +def test_tile_crossing_polygon_is_not_duplicated(written: tuple[Path, dict]) -> None: + """Its outline spans two tiles but it must appear exactly once, in its centroid's tile.""" + out, _ = written + table = pq.read_table(out) + codes = table[CELL_CODE_COLUMN].to_pylist() + crossing = list(SHAPES_SPEC).index("cell-crossing") + assert codes.count(crossing) == 1 + assert table.num_rows == len(SHAPES_SPEC) + + +def test_every_cell_appears_exactly_once(written: tuple[Path, dict]) -> None: + out, _ = written + codes = pq.read_table(out)[CELL_CODE_COLUMN].to_pylist() + assert sorted(codes) == list(range(len(SHAPES_SPEC))) + + +def test_empty_tile_is_zero_rows(written: tuple[Path, dict]) -> None: + out, _ = written + # tiles 3 and 4 hold no cells + f = pq.ParquetFile(out) + for tile_id in (3, 4): + assert f.metadata.row_group(tile_id).num_rows == 0 + + +# -- canonical preservation --------------------------------------------------- + + +def test_canonical_geometry_is_unchanged(written: tuple[Path, dict], shapes: gpd.GeoDataFrame) -> None: + out, _ = written + back = gpd.read_parquet(out) + for name, (geom, _) in SHAPES_SPEC.items(): + assert back.loc[name].geometry.equals(geom), name + + +def test_output_is_still_valid_geoparquet(written: tuple[Path, dict]) -> None: + out, _ = written + assert b"geo" in (pq.ParquetFile(out).schema_arrow.metadata or {}) + back = gpd.read_parquet(out) + assert isinstance(back, gpd.GeoDataFrame) + assert back.geometry.name == "geometry" + + +def test_non_geometry_columns_survive(written: tuple[Path, dict], shapes: gpd.GeoDataFrame) -> None: + out, _ = written + back = gpd.read_parquet(out) + for name in SHAPES_SPEC: + assert back.loc[name, "area"] == pytest.approx(shapes.loc[name, "area"]) + + +# -- display geometry --------------------------------------------------------- + + +def test_display_geometry_has_the_nested_layout(written: tuple[Path, dict]) -> None: + """polygon -> rings -> interleaved uint32 pairs, as get_polygon_data.js walks it.""" + out, _ = written + t = pq.ParquetFile(out).schema_arrow.field(GEOMETRY_COLUMN).type + assert pa.types.is_list(t) # polygon level + assert pa.types.is_list(t.value_type) # ring level + assert t.value_type.value_type == pa.list_(pa.uint32(), 2) # interleaved vertices + + +def test_display_geometry_offsets_resolve_like_the_js_reader(written: tuple[Path, dict]) -> None: + """Mirror of getPolygonDataFromChunk: polygon offset -> ring offset -> coord index.""" + out, _ = written + col = pq.read_table(out)[GEOMETRY_COLUMN].combine_chunks() + polygon_offsets = col.offsets.to_numpy() + rings = col.values + ring_offsets = rings.offsets.to_numpy() + flat = rings.values.values.to_numpy(zero_copy_only=False) + + start_indices = ring_offsets[polygon_offsets] + assert len(start_indices) == len(col) + 1 + # First polygon's first vertex, read the way the browser would. + first = flat[2 * start_indices[0] : 2 * start_indices[0] + 2] + assert first.tolist() == col.to_pylist()[0][0][0] + + +def test_display_geometry_is_exterior_ring_only(written: tuple[Path, dict]) -> None: + out, _ = written + for polygon in pq.read_table(out)[GEOMETRY_COLUMN].to_pylist(): + assert len(polygon) == 1, "expected exactly one ring per display polygon" + + +def test_display_vertices_match_the_transform(written: tuple[Path, dict]) -> None: + out, _ = written + table = pq.read_table(out) + codes = table[CELL_CODE_COLUMN].to_pylist() + geoms = table[GEOMETRY_COLUMN].to_pylist() + names = list(SHAPES_SPEC) + for code, poly in zip(codes, geoms): + canonical = SHAPES_SPEC[names[code]][0] + expected = [[int(round(x * 2)), int(round(y * 2))] for x, y in canonical.exterior.coords] + assert [list(v) for v in poly[0]] == expected + + +def test_multipolygon_reduces_to_largest_part(tmp_path: Path) -> None: + from spatialdata.models import ShapesModel + + big, small = _square(2.5, 2.5, 2.0), _square(8.0, 13.0, 0.5) + gdf = gpd.GeoDataFrame(geometry=[MultiPolygon([big, small])], index=["multi"]) + out = tmp_path / "m.parquet" + write_shapes_regular_grid(ShapesModel.parse(gdf), out, grid=GRID, display_transform=XFORM) + poly = pq.read_table(out)[GEOMETRY_COLUMN].to_pylist()[0] + got = {tuple(v) for v in poly[0]} + assert got == {(int(round(x * 2)), int(round(y * 2))) for x, y in big.exterior.coords} + + +# -- cell codes --------------------------------------------------------------- + + +def test_cell_codes_follow_the_table_order(tmp_path: Path, shapes: gpd.GeoDataFrame) -> None: + """cell_code must index the annotating table, not the shapes' own row order.""" + reversed_index = list(SHAPES_SPEC)[::-1] + out = tmp_path / "c.parquet" + write_shapes_regular_grid(shapes, out, grid=GRID, display_transform=XFORM, cell_index=reversed_index) + back = gpd.read_parquet(out) + for name in SHAPES_SPEC: + assert back.loc[name, CELL_CODE_COLUMN] == reversed_index.index(name) + + +def test_cell_missing_from_index_is_reported(tmp_path: Path, shapes: gpd.GeoDataFrame) -> None: + with pytest.raises(ValueError, match="absent from cell_index"): + write_shapes_regular_grid( + shapes, tmp_path / "x.parquet", grid=GRID, display_transform=XFORM, cell_index=["cell-a"] + ) + + +def test_overwrite_guard(written: tuple[Path, dict], shapes: gpd.GeoDataFrame) -> None: + out, _ = written + with pytest.raises(FileExistsError): + write_shapes_regular_grid(shapes, out, grid=GRID, display_transform=XFORM) From 8e0f8a0c86a5b71a84bcf8aec5a6f83d6e56786a Mon Sep 17 00:00:00 2001 From: Nicolas Fernandez Date: Sun, 6 Sep 2026 14:20:46 -0400 Subject: [PATCH 04/17] feat(experimental): add gene-major cell-by-gene writer One row group per gene, so selecting a gene fetches a single row group instead of touching any transcript data. Schema matches Celldega's existing CBG reader exactly: cell_id / expression / gene, with gene_to_row_group and num_genes in the parquet schema metadata. Deliberate difference from Celldega's own writer: every catalog gene gets a row group in catalog order, including all-zero genes, so that 'row group index == feature_code' holds and one integer addresses both a transcript's gene and its expression vector. gene_to_row_group is still written for clients that look it up rather than assume it. Explicit stored zeros are dropped -- a sparse matrix can carry them and they are not expression. Validated against Xenium pancreas: 377 genes, 2,607,168 values, 16.2 MB, 0.2s; INS/GCG/ACTA2 cell sets and values match the AnnData table exactly. Co-Authored-By: Claude Opus 5 (1M context) --- .../experimental/cbg_parquet.py | 181 +++++++++++++++++ tests/test_cbg_parquet.py | 190 ++++++++++++++++++ 2 files changed, 371 insertions(+) create mode 100644 src/spatialdata_io/experimental/cbg_parquet.py create mode 100644 tests/test_cbg_parquet.py diff --git a/src/spatialdata_io/experimental/cbg_parquet.py b/src/spatialdata_io/experimental/cbg_parquet.py new file mode 100644 index 00000000..f1779b1a --- /dev/null +++ b/src/spatialdata_io/experimental/cbg_parquet.py @@ -0,0 +1,181 @@ +"""Write a gene-major cell-by-gene matrix as one row group per gene. + +This is what lets the overview render cheaply: selecting a gene fetches a single row +group holding that gene's non-zero cells, instead of touching any transcript data. + +The schema matches Celldega's existing CBG reader exactly -- columns ``cell_id`` +(the integer cell code, not a string barcode), ``expression``, ``gene``, plus +``gene_to_row_group`` / ``num_genes`` in the Parquet schema metadata. + +One deliberate difference from Celldega's own writer: every gene in the catalog gets a +row group, in catalog order, even if it has no non-zero cells. That preserves the +invariant ``feature_code == cbg row group``, so a transcript's feature code and its +expression vector are addressed by the same integer. ``gene_to_row_group`` is still +written, so a client that looks the mapping up rather than assuming it works unchanged. +""" + +from __future__ import annotations + +import json +import shutil +from pathlib import Path +from typing import Any + +import numpy as np +import pyarrow as pa +import pyarrow.parquet as pq +import scipy.sparse as sp + +from spatialdata_io.experimental.feature_catalog import FeatureCatalog +from spatialdata_io.experimental.regular_grid import DEFAULT_MAX_ROW_GROUPS_PER_FILE + +__all__ = ["write_cbg_row_groups"] + + +def write_cbg_row_groups( + table: Any, + output_dir: str | Path, + *, + catalog: FeatureCatalog, + cell_codes: dict[str, int] | None = None, + layer: str | None = None, + max_row_groups_per_file: int = DEFAULT_MAX_ROW_GROUPS_PER_FILE, + compression: str = "snappy", + overwrite: bool = False, +) -> dict[str, Any]: + """Write an AnnData table as gene-major CBG row groups. + + Parameters + ---------- + table + The annotating :class:`anndata.AnnData` table. + output_dir + Directory to write the chunk files into. Written atomically. + catalog + Feature catalog; its gene order defines the row-group order. Every gene in the + catalog must be present in ``table.var_names``. + cell_codes + Mapping from ``obs`` name to integer cell code. Defaults to positional order, + which is what :func:`write_shapes_regular_grid` uses by default. + layer + Optional ``table.layers`` key to read instead of ``table.X``. + max_row_groups_per_file + Genes per chunk file. + compression + Parquet compression codec. + overwrite + Replace ``output_dir`` if it exists. + + Returns + ------- + The manifest fragment describing the written files. + """ + output_dir = Path(output_dir) + if output_dir.exists() and not overwrite: + raise FileExistsError(f"{output_dir} exists; pass overwrite=True to replace it") + + var_names = list(table.var_names) + positions = {g: i for i, g in enumerate(var_names)} + missing = [g for g in catalog.genes if g not in positions] + if missing: + raise ValueError( + f"{len(missing)} catalog gene(s) are absent from table.var_names " + f"(e.g. {missing[:3]}); the CBG would not cover every gene code." + ) + + if cell_codes is None: + codes = np.arange(table.n_obs, dtype=np.uint32) + else: + unknown = [k for k in table.obs_names if k not in cell_codes] + if unknown: + raise ValueError(f"{len(unknown)} table cell(s) have no cell_code (e.g. {unknown[:3]})") + codes = np.fromiter((cell_codes[k] for k in table.obs_names), dtype=np.uint32, count=table.n_obs) + + matrix = table.layers[layer] if layer is not None else table.X + # CSC gives O(1) access to a single gene's column, which is the whole point here. + csc = matrix.tocsc() if sp.issparse(matrix) else sp.csc_matrix(np.asarray(matrix)) + csc.sort_indices() + + schema = pa.schema( + [ + pa.field("cell_id", pa.uint32()), + pa.field("expression", pa.float32()), + pa.field("gene", pa.string()), + ] + ) + + gene_to_row_group = {gene: i for i, gene in enumerate(catalog.genes)} + n_genes = len(catalog.genes) + n_files = -(-n_genes // max_row_groups_per_file) + width = len(str(n_files - 1)) if n_files > 1 else 1 + filenames = [f"chunk_{i:0{width}d}.parquet" for i in range(n_files)] + + schema = schema.with_metadata( + { + b"gene_to_row_group": json.dumps(gene_to_row_group).encode(), + b"storage_mode": b"row_groups_cbg_chunked", + b"num_genes": str(n_genes).encode(), + b"max_row_groups_per_file": str(max_row_groups_per_file).encode(), + b"profile": b"celldega_regular_grid_v1", + } + ) + + staging = output_dir.with_name(output_dir.name + ".tmp") + if staging.exists(): + shutil.rmtree(staging) + staging.mkdir(parents=True) + + n_values = 0 + try: + writer: pq.ParquetWriter | None = None + current_file = -1 + for row_group, gene in enumerate(catalog.genes): + file_index = row_group // max_row_groups_per_file + if file_index != current_file: + if writer is not None: + writer.close() + writer = pq.ParquetWriter( + staging / filenames[file_index], schema, compression=compression, write_statistics=False + ) + current_file = file_index + + col = positions[gene] + start, end = csc.indptr[col], csc.indptr[col + 1] + cells = codes[csc.indices[start:end]] + values = csc.data[start:end] + # Explicit zeros can survive in a sparse matrix; the CBG stores non-zeros only. + keep = values != 0 + cells, values = cells[keep], values[keep] + n_values += len(values) + + assert writer is not None + writer.write_table( + pa.table( + { + "cell_id": pa.array(cells, type=pa.uint32()), + "expression": pa.array(values, type=pa.float32()), + "gene": pa.array([gene] * len(values), type=pa.string()), + }, + schema=schema, + ) + ) + if writer is not None: + writer.close() + except BaseException: + shutil.rmtree(staging, ignore_errors=True) + raise + + if output_dir.exists(): + shutil.rmtree(output_dir) + staging.rename(output_dir) + + return { + "directory": output_dir.name, + "files": filenames, + "max_row_groups_per_file": max_row_groups_per_file, + "total_row_groups": n_genes, + "num_genes": n_genes, + "gene_to_row_group": gene_to_row_group, + "n_values": n_values, + "row_group_equals_feature_code": True, + } diff --git a/tests/test_cbg_parquet.py b/tests/test_cbg_parquet.py new file mode 100644 index 00000000..6eb07aec --- /dev/null +++ b/tests/test_cbg_parquet.py @@ -0,0 +1,190 @@ +"""Tests for the gene-major cell-by-gene writer.""" + +from __future__ import annotations + +import json +from pathlib import Path + +import numpy as np +import pandas as pd +import pyarrow.parquet as pq +import pytest +import scipy.sparse as sp +from anndata import AnnData + +from spatialdata_io.experimental.cbg_parquet import write_cbg_row_groups +from spatialdata_io.experimental.feature_catalog import FeatureCatalog + +GENES = ["GENEA", "GENEB", "GENEC"] +CELLS = ["cell-0", "cell-1", "cell-2", "cell-3"] + +#: rows = cells, cols = genes. GENEC is deliberately all-zero. +DENSE = np.array( + [ + [5.0, 0.0, 0.0], + [0.0, 2.0, 0.0], + [3.0, 7.0, 0.0], + [0.0, 0.0, 0.0], + ], + dtype=np.float32, +) + + +@pytest.fixture +def table() -> AnnData: + return AnnData( + X=sp.csr_matrix(DENSE), + obs=pd.DataFrame(index=CELLS), + var=pd.DataFrame(index=GENES), + ) + + +@pytest.fixture +def catalog() -> FeatureCatalog: + return FeatureCatalog.from_features_and_table([*GENES, "NegControlProbe_0001"], GENES) + + +@pytest.fixture +def written(tmp_path: Path, table: AnnData, catalog: FeatureCatalog) -> tuple[Path, dict]: + out = tmp_path / "cbg" + manifest = write_cbg_row_groups(table, out, catalog=catalog) + return out, manifest + + +def _row_group_for(directory: Path, manifest: dict, gene: str): + rg = manifest["gene_to_row_group"][gene] + file_index, local = divmod(rg, manifest["max_row_groups_per_file"]) + return pq.ParquetFile(directory / manifest["files"][file_index]).read_row_group(local) + + +# -- layout ------------------------------------------------------------------- + + +def test_one_row_group_per_gene(written: tuple[Path, dict]) -> None: + directory, manifest = written + total = sum(pq.ParquetFile(directory / f).metadata.num_row_groups for f in manifest["files"]) + assert total == len(GENES) == manifest["num_genes"] + + +def test_row_group_index_equals_feature_code(written: tuple[Path, dict], catalog: FeatureCatalog) -> None: + """The invariant that lets one integer address both a transcript's gene and its CBG vector.""" + _, manifest = written + for gene, rg in manifest["gene_to_row_group"].items(): + assert catalog.names.index(gene) == rg + assert manifest["row_group_equals_feature_code"] is True + + +def test_controls_get_no_row_group(written: tuple[Path, dict]) -> None: + _, manifest = written + assert "NegControlProbe_0001" not in manifest["gene_to_row_group"] + assert manifest["num_genes"] == len(GENES) + + +def test_schema_matches_celldega_reader(written: tuple[Path, dict]) -> None: + directory, manifest = written + f = pq.ParquetFile(directory / manifest["files"][0]) + assert f.schema_arrow.names == ["cell_id", "expression", "gene"] + meta = f.schema_arrow.metadata + assert json.loads(meta[b"gene_to_row_group"]) == manifest["gene_to_row_group"] + assert meta[b"storage_mode"] == b"row_groups_cbg_chunked" + assert int(meta[b"num_genes"]) == len(GENES) + + +# -- values ------------------------------------------------------------------- + + +def test_sparse_values_match_the_source_matrix(written: tuple[Path, dict]) -> None: + directory, manifest = written + for col, gene in enumerate(GENES): + rg = _row_group_for(directory, manifest, gene).to_pandas() + expected = {i: DENSE[i, col] for i in range(len(CELLS)) if DENSE[i, col] != 0} + assert dict(zip(rg["cell_id"], rg["expression"])) == expected + assert set(rg["gene"]) <= {gene} + + +def test_zero_values_are_omitted(written: tuple[Path, dict]) -> None: + directory, manifest = written + for gene in GENES: + rg = _row_group_for(directory, manifest, gene) + assert all(v != 0 for v in rg["expression"].to_pylist()) + + +def test_all_zero_gene_is_an_empty_row_group(written: tuple[Path, dict]) -> None: + """GENEC still occupies its slot so the feature_code invariant holds.""" + directory, manifest = written + assert "GENEC" in manifest["gene_to_row_group"] + assert _row_group_for(directory, manifest, "GENEC").num_rows == 0 + + +def test_explicit_stored_zeros_are_dropped(tmp_path: Path, catalog: FeatureCatalog) -> None: + """A sparse matrix can carry stored zeros; they are not expression.""" + m = sp.csr_matrix(DENSE) + m[3, 0] = 0 # creates an explicit stored zero + adata = AnnData(X=m, obs=pd.DataFrame(index=CELLS), var=pd.DataFrame(index=GENES)) + out = tmp_path / "cbg" + manifest = write_cbg_row_groups(adata, out, catalog=catalog) + rg = _row_group_for(out, manifest, "GENEA").to_pandas() + assert 3 not in set(rg["cell_id"]) + + +def test_dense_matrix_is_supported(tmp_path: Path, catalog: FeatureCatalog) -> None: + adata = AnnData(X=DENSE.copy(), obs=pd.DataFrame(index=CELLS), var=pd.DataFrame(index=GENES)) + out = tmp_path / "cbg" + manifest = write_cbg_row_groups(adata, out, catalog=catalog) + rg = _row_group_for(out, manifest, "GENEB").to_pandas() + assert dict(zip(rg["cell_id"], rg["expression"])) == {1: 2.0, 2: 7.0} + + +def test_layer_can_be_selected(tmp_path: Path, table: AnnData, catalog: FeatureCatalog) -> None: + table.layers["scaled"] = sp.csr_matrix(DENSE * 10) + out = tmp_path / "cbg" + manifest = write_cbg_row_groups(table, out, catalog=catalog, layer="scaled") + rg = _row_group_for(out, manifest, "GENEA").to_pandas() + assert dict(zip(rg["cell_id"], rg["expression"])) == {0: 50.0, 2: 30.0} + + +# -- cell codes --------------------------------------------------------------- + + +def test_cell_codes_default_to_table_order(written: tuple[Path, dict]) -> None: + directory, manifest = written + rg = _row_group_for(directory, manifest, "GENEA").to_pandas() + assert sorted(rg["cell_id"]) == [0, 2] + + +def test_explicit_cell_codes_are_honoured(tmp_path: Path, table: AnnData, catalog: FeatureCatalog) -> None: + codes = {name: i for i, name in enumerate(reversed(CELLS))} + out = tmp_path / "cbg" + manifest = write_cbg_row_groups(table, out, catalog=catalog, cell_codes=codes) + rg = _row_group_for(out, manifest, "GENEA").to_pandas() + assert sorted(rg["cell_id"]) == sorted([codes["cell-0"], codes["cell-2"]]) + + +def test_missing_cell_code_is_reported(tmp_path: Path, table: AnnData, catalog: FeatureCatalog) -> None: + with pytest.raises(ValueError, match="no cell_code"): + write_cbg_row_groups(table, tmp_path / "cbg", catalog=catalog, cell_codes={"cell-0": 0}) + + +def test_gene_missing_from_table_is_reported(tmp_path: Path, table: AnnData) -> None: + bad = FeatureCatalog(names=("GENEA", "GHOST"), n_genes=2) + with pytest.raises(ValueError, match="absent from table.var_names"): + write_cbg_row_groups(table, tmp_path / "cbg", catalog=bad) + + +# -- chunking ----------------------------------------------------------------- + + +def test_multi_file_chunking(tmp_path: Path, table: AnnData, catalog: FeatureCatalog) -> None: + out = tmp_path / "cbg" + manifest = write_cbg_row_groups(table, out, catalog=catalog, max_row_groups_per_file=2) + assert manifest["files"] == ["chunk_0.parquet", "chunk_1.parquet"] + assert [pq.ParquetFile(out / f).metadata.num_row_groups for f in manifest["files"]] == [2, 1] + # the mapping must still resolve through the file split + rg = _row_group_for(out, manifest, "GENEC") + assert rg.num_rows == 0 + + +def test_overwrite_guard(written: tuple[Path, dict], table: AnnData, catalog: FeatureCatalog) -> None: + directory, _ = written + with pytest.raises(FileExistsError): + write_cbg_row_groups(table, directory, catalog=catalog) From 410d3b8671e8bf6ec8ac666ccaef0055d497566d Mon Sep 17 00:00:00 2001 From: Nicolas Fernandez Date: Sun, 6 Sep 2026 14:30:51 -0400 Subject: [PATCH 05/17] perf(experimental): default to zstd compression Measured on Xenium pancreas (8.07M transcripts, 250px tiles), total overhead of the two render columns plus row-group fragmentation, against the 249.1 MB SpatialData default write: snappy 345.1 MB +38.5% 7.0s zstd 261.4 MB +4.9% 7.3s brotli 250.4 MB +0.5% 50.1s snappy barely compresses display_xy at all; zstd removes most of the overhead for no meaningful write cost, and brotli is 7x slower for another 4 points. Verified parquet-wasm 0.7.1 (celldega's pinned version) can read zstd: ZSTD = 5 is in its Compression enum and the zstd shims are compiled into the wasm bundle. This matters because an unsupported codec would fail only in the browser. Co-Authored-By: Claude Opus 5 (1M context) --- src/spatialdata_io/experimental/cbg_parquet.py | 2 +- src/spatialdata_io/experimental/points_parquet.py | 2 +- src/spatialdata_io/experimental/shapes_parquet.py | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/src/spatialdata_io/experimental/cbg_parquet.py b/src/spatialdata_io/experimental/cbg_parquet.py index f1779b1a..a717674a 100644 --- a/src/spatialdata_io/experimental/cbg_parquet.py +++ b/src/spatialdata_io/experimental/cbg_parquet.py @@ -40,7 +40,7 @@ def write_cbg_row_groups( cell_codes: dict[str, int] | None = None, layer: str | None = None, max_row_groups_per_file: int = DEFAULT_MAX_ROW_GROUPS_PER_FILE, - compression: str = "snappy", + compression: str = "zstd", overwrite: bool = False, ) -> dict[str, Any]: """Write an AnnData table as gene-major CBG row groups. diff --git a/src/spatialdata_io/experimental/points_parquet.py b/src/spatialdata_io/experimental/points_parquet.py index 65c30e90..4bb11d2a 100644 --- a/src/spatialdata_io/experimental/points_parquet.py +++ b/src/spatialdata_io/experimental/points_parquet.py @@ -139,7 +139,7 @@ def write_points_regular_grid( feature_key: str = "feature_name", tile_size_px: float = 250.0, max_row_groups_per_file: int = DEFAULT_MAX_ROW_GROUPS_PER_FILE, - compression: str = "snappy", + compression: str = "zstd", overwrite: bool = False, ) -> dict[str, Any]: """Write a Points element as regular-grid row groups. diff --git a/src/spatialdata_io/experimental/shapes_parquet.py b/src/spatialdata_io/experimental/shapes_parquet.py index 208f6953..d252c6ee 100644 --- a/src/spatialdata_io/experimental/shapes_parquet.py +++ b/src/spatialdata_io/experimental/shapes_parquet.py @@ -124,7 +124,7 @@ def write_shapes_regular_grid( coordinate_system: str = "global", cell_index: Any | None = None, max_row_groups_per_file: int = DEFAULT_MAX_ROW_GROUPS_PER_FILE, - compression: str = "snappy", + compression: str = "zstd", overwrite: bool = False, ) -> dict[str, Any]: """Write a Shapes element as regular-grid row groups. From 37c9ea0e1f2b811564e955f2705d2e1ca41932b9 Mon Sep 17 00:00:00 2001 From: Nicolas Fernandez Date: Sun, 6 Sep 2026 14:40:52 -0400 Subject: [PATCH 06/17] feat(experimental): add profile manifest and opt-in tiling entry points add_spatial_tiling() tiles an existing store in place; xenium_spatially_tiled() goes from raw Xenium to a tiled store in one call. Both are single calls; the one-shot path is internally read -> write -> tile because the tiling rewrites written parquet. The manifest reuses Celldega's landscape_parameters.json keys (use_row_groups, tile_grid, row_group_files, technology, image_info) so its reader consumes it unchanged, and declares paths relative to the profile directory (../../points/transcripts/points.parquet) so Celldega can be pointed at that directory as base_url with no reader change. validate_manifest() fails loudly on the mistakes that would otherwise surface as a blank viewport: row-group counts disagreeing with the tile grid, file counts that cannot hold the declared row groups, CBG mappings pointing past the end, and declared files that do not exist. The re-runnability test caught the same idempotency bug in the shapes writer that was previously fixed for points: re-tiling appended duplicate display_geometry / cell_code columns rather than replacing them. Verified on Xenium pancreas: 15.0s to tile, 7,535 row groups across 19 files, read_zarr still returns 8,073,840 points / 140,702 shapes / (140702, 377) table, and canonical geometry is byte-identical after the reorder (including the 159 pre-existing invalid polygons, which are 10x source data, not corruption). Co-Authored-By: Claude Opus 5 (1M context) --- src/spatialdata_io/experimental/manifest.py | 173 ++++++++++++ .../experimental/shapes_parquet.py | 6 +- .../experimental/tiled_access.py | 263 ++++++++++++++++++ tests/test_tiled_access.py | 240 ++++++++++++++++ 4 files changed, 681 insertions(+), 1 deletion(-) create mode 100644 src/spatialdata_io/experimental/manifest.py create mode 100644 src/spatialdata_io/experimental/tiled_access.py create mode 100644 tests/test_tiled_access.py diff --git a/src/spatialdata_io/experimental/manifest.py b/src/spatialdata_io/experimental/manifest.py new file mode 100644 index 00000000..58f4c9af --- /dev/null +++ b/src/spatialdata_io/experimental/manifest.py @@ -0,0 +1,173 @@ +"""The visualization profile manifest. + +The manifest is what turns a pile of Parquet files into a discoverable profile: it tells a +client the tile geometry, which files hold which row groups, and which columns to project. +Without it a client would have to infer the layout, which is exactly what this profile +exists to avoid. + +Key names follow Celldega's existing ``landscape_parameters.json`` so that its reader can +consume the manifest unchanged (``use_row_groups``, ``tile_grid``, ``row_group_files``, +``technology``, ``image_info``). Profile-specific additions live under ``profile`` and in +each entry's column names, so a client that does not understand them still finds the keys +it expects. +""" + +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any + +from spatialdata_io.experimental.regular_grid import RegularGrid + +__all__ = [ + "PROFILE_NAME", + "PROFILE_VERSION", + "MANIFEST_FILENAME", + "build_manifest", + "validate_manifest", + "write_manifest", +] + +PROFILE_NAME = "celldega_regular_grid_v1" +PROFILE_VERSION = "0.1.0" +MANIFEST_FILENAME = "landscape_parameters.json" + + +def build_manifest( + *, + grid: RegularGrid, + technology: str = "Xenium", + transcripts: dict[str, Any] | None = None, + cell_segmentation: dict[str, Any] | None = None, + cbg: dict[str, Any] | None = None, + images: dict[str, Any] | None = None, + image_info: list[dict[str, Any]] | None = None, + source: dict[str, Any] | None = None, +) -> dict[str, Any]: + """Assemble the profile manifest from the fragments returned by each writer. + + Parameters + ---------- + grid + The tile grid shared by transcripts and shapes. + technology + Celldega technology string, controlling viewer behaviour such as whether an image + layer exists. + transcripts, cell_segmentation, cbg, images + Fragments returned by the corresponding writers. Omitted entries are simply absent, + and Celldega skips initialising a reader for them. + image_info + Celldega image descriptors. An empty list is valid and means "no image layer". + source + Provenance recorded for invalidation: which store and elements this was built from. + + Returns + ------- + The manifest as a plain dict. + """ + row_group_files: dict[str, Any] = {} + if transcripts is not None: + row_group_files["transcripts"] = transcripts + if cell_segmentation is not None: + row_group_files["cell_segmentation"] = cell_segmentation + if cbg is not None: + row_group_files["cbg"] = cbg + row_group_files["images"] = images or {} + + manifest: dict[str, Any] = { + # -- keys Celldega's existing reader consumes ------------------------- + "technology": technology, + "use_row_groups": True, + "tile_grid": grid.to_manifest_dict(), + "row_group_files": row_group_files, + "image_info": image_info or [], + # -- profile identification ------------------------------------------- + "profile": PROFILE_NAME, + "profile_version": PROFILE_VERSION, + } + if source is not None: + manifest["source"] = source + return manifest + + +def validate_manifest(manifest: dict[str, Any], base_path: str | Path | None = None) -> None: + """Check a manifest is self-consistent, and that its files exist when ``base_path`` is given. + + Raises + ------ + ValueError + With a message naming the specific problem. Failing here is much cheaper than + failing as a blank viewport in a browser. + """ + for key in ("technology", "use_row_groups", "tile_grid", "row_group_files"): + if key not in manifest: + raise ValueError(f"manifest is missing required key {key!r}") + + grid = RegularGrid.from_manifest_dict(manifest["tile_grid"]) + files = manifest["row_group_files"] + base = Path(base_path) if base_path is not None else None + + for name in ("transcripts", "cell_segmentation"): + entry = files.get(name) + if entry is None: + continue + expected = grid.num_tiles + if entry.get("total_row_groups") != expected: + raise ValueError( + f"{name}: total_row_groups={entry.get('total_row_groups')} does not match the " + f"tile grid ({grid.num_tiles_x} x {grid.num_tiles_y} = {expected} tiles). " + f"The row-group index formula would address the wrong tiles." + ) + _check_paths(name, entry, base) + + cbg = files.get("cbg") + if cbg is not None: + mapping = cbg.get("gene_to_row_group", {}) + if not mapping: + raise ValueError("cbg: gene_to_row_group is empty; no gene could be selected") + if cbg.get("total_row_groups") != len(mapping): + raise ValueError( + f"cbg: total_row_groups={cbg.get('total_row_groups')} does not match " + f"{len(mapping)} entries in gene_to_row_group" + ) + max_rg = cbg.get("max_row_groups_per_file", 1) + n_files = len(cbg.get("files", [])) + if mapping and max(mapping.values()) >= n_files * max_rg: + raise ValueError( + f"cbg: gene_to_row_group references row group {max(mapping.values())} but only " + f"{n_files} file(s) x {max_rg} row groups are listed" + ) + _check_paths("cbg", cbg, base) + + +def _check_paths(name: str, entry: dict[str, Any], base: Path | None) -> None: + """Verify the declared files exist, and that the entry names them coherently.""" + if "files" in entry: + if "directory" not in entry: + raise ValueError(f"{name}: entry lists 'files' but no 'directory'") + n_files = len(entry["files"]) + expected = -(-entry["total_row_groups"] // entry["max_row_groups_per_file"]) + if n_files != expected: + raise ValueError( + f"{name}: lists {n_files} file(s) but {entry['total_row_groups']} row groups at " + f"{entry['max_row_groups_per_file']} per file need {expected}" + ) + if base is not None: + for f in entry["files"]: + if not (base / entry["directory"] / f).exists(): + raise ValueError(f"{name}: declared file {entry['directory']}/{f} does not exist") + elif "path" in entry: + if base is not None and not (base / entry["path"]).exists(): + raise ValueError(f"{name}: declared path {entry['path']} does not exist") + else: + raise ValueError(f"{name}: entry has neither 'files' nor 'path'") + + +def write_manifest(manifest: dict[str, Any], directory: str | Path) -> Path: + """Write the manifest as ``landscape_parameters.json`` and return its path.""" + directory = Path(directory) + directory.mkdir(parents=True, exist_ok=True) + path = directory / MANIFEST_FILENAME + path.write_text(json.dumps(manifest, indent=2, sort_keys=False)) + return path diff --git a/src/spatialdata_io/experimental/shapes_parquet.py b/src/spatialdata_io/experimental/shapes_parquet.py index d252c6ee..b5a2122e 100644 --- a/src/spatialdata_io/experimental/shapes_parquet.py +++ b/src/spatialdata_io/experimental/shapes_parquet.py @@ -163,7 +163,11 @@ def write_shapes_regular_grid( transform = display_transform or DisplayTransform.from_element(shapes, coordinate_system) display, cx, cy = _display_geometry_array(shapes.geometry, transform) - table = _canonical_geoparquet_table(shapes) + # Re-tiling an already-tiled element must replace the render columns, not append + # duplicates: a duplicated name makes projected reads fail outright, and the stale + # copy is also mistyped because a pandas round-trip degrades fixed_size_list to list. + stale = [c for c in (GEOMETRY_COLUMN, CELL_CODE_COLUMN) if c in shapes.columns] + table = _canonical_geoparquet_table(shapes.drop(columns=stale) if stale else shapes) if cell_index is None: codes = np.arange(len(shapes), dtype=np.uint32) diff --git a/src/spatialdata_io/experimental/tiled_access.py b/src/spatialdata_io/experimental/tiled_access.py new file mode 100644 index 00000000..53fb1fd8 --- /dev/null +++ b/src/spatialdata_io/experimental/tiled_access.py @@ -0,0 +1,263 @@ +"""Opt-in spatial tiling for a SpatialData store. + +Two entry points, both a single call: + +:func:`add_spatial_tiling` + Add the profile to a store that already exists. +:func:`xenium_spatially_tiled` + Read raw Xenium and write a tiled store in one go. + +The one-shot path is internally ``read -> write -> tile`` because the tiling rewrites +*written* Parquet. When the installed SpatialData exposes the ``points_writer`` hook the +points element is written in its tiled form directly, avoiding writing it twice; otherwise +it falls back to writing normally and rewriting, which produces an identical store. + +Everything here is additive and opt-in. A store that has been tiled is still an ordinary +SpatialData store: :func:`spatialdata.read_zarr` works unchanged, the canonical columns and +geometries are untouched, and a client that does not know about the profile simply ignores +the extra columns and the manifest. +""" + +from __future__ import annotations + +import inspect +import shutil +from pathlib import Path +from typing import Any + +import numpy as np + +from spatialdata_io.experimental.cbg_parquet import write_cbg_row_groups +from spatialdata_io.experimental.feature_catalog import FeatureCatalog +from spatialdata_io.experimental.manifest import ( + PROFILE_NAME, + build_manifest, + validate_manifest, + write_manifest, +) +from spatialdata_io.experimental.points_parquet import ( + DisplayTransform, + write_points_regular_grid, +) +from spatialdata_io.experimental.regular_grid import ( + DEFAULT_MAX_ROW_GROUPS_PER_FILE, + RegularGrid, +) +from spatialdata_io.experimental.shapes_parquet import write_shapes_regular_grid + +__all__ = ["add_spatial_tiling", "xenium_spatially_tiled", "supports_points_writer_hook"] + +#: Directory inside the store holding derived (non-canonical) profile assets. +PROFILE_DIR = "visualization" + + +def supports_points_writer_hook() -> bool: + """Whether the installed SpatialData accepts a ``points_writer`` in ``write()``.""" + from spatialdata import SpatialData + + return "points_writer" in inspect.signature(SpatialData.write).parameters + + +def _grid_for(points: Any, transform: DisplayTransform, tile_size_px: float) -> RegularGrid: + """Derive the grid covering the points element in display pixel space.""" + x = points["x"].max().compute() if hasattr(points["x"].max(), "compute") else points["x"].max() + y = points["y"].max().compute() if hasattr(points["y"].max(), "compute") else points["y"].max() + px, py = transform.apply(np.array([float(x)]), np.array([float(y)])) + return RegularGrid.from_bounds(0, 0, float(np.rint(px[0])), float(np.rint(py[0])), tile_size_px) + + +def add_spatial_tiling( + store: str | Path, + *, + points_element: str = "transcripts", + shapes_element: str | None = "cell_boundaries", + table_element: str | None = "table", + coordinate_system: str = "global", + tile_size_px: float = 250.0, + max_row_groups_per_file: int = DEFAULT_MAX_ROW_GROUPS_PER_FILE, + feature_key: str = "feature_name", + technology: str = "Xenium", + include_cbg: bool = True, + compression: str = "zstd", +) -> dict[str, Any]: + """Add the regular-grid visualization profile to an existing SpatialData store. + + Re-runnable: running it again replaces the render columns and derived assets rather + than duplicating them. + + Parameters + ---------- + store + Path to the ``.zarr`` store to tile, modified in place. + points_element + Name of the Points element holding transcripts. + shapes_element + Name of the Shapes element holding cell boundaries, or ``None`` to skip. + table_element + Name of the annotating table, used for the gene order and the CBG. ``None`` skips + the CBG and derives feature codes from the observed features alone. + coordinate_system + Coordinate system defining display pixel space. + tile_size_px + Tile edge length in display pixels. The default of 250 gives roughly 20 cells per + tile on Xenium-density tissue, which is the granularity the viewer fetches at. + max_row_groups_per_file + Row groups per chunk file. + feature_key + Column in the points element holding the feature name. + technology + Celldega technology string recorded in the manifest. + include_cbg + Whether to write the gene-major cell-by-gene files. + compression + Parquet compression codec. + + Returns + ------- + The profile manifest. + """ + import spatialdata + + store = Path(store) + sdata = spatialdata.read_zarr(store) + + if points_element not in sdata.points: + raise ValueError(f"points element {points_element!r} not found; have {list(sdata.points)}") + points = sdata.points[points_element] + table = sdata.tables[table_element] if table_element else None + + transform = DisplayTransform.from_element(points, coordinate_system) + grid = _grid_for(points, transform, tile_size_px) + + if table is not None: + catalog = FeatureCatalog.from_points_and_table(points, table, feature_key=feature_key) + else: + observed = points[feature_key] + names = sorted(observed.cat.as_known().cat.categories) if hasattr(observed, "cat") else [] + catalog = FeatureCatalog(names=tuple(names), n_genes=len(names)) + + profile_dir = store / PROFILE_DIR / PROFILE_NAME + profile_dir.mkdir(parents=True, exist_ok=True) + + transcripts = write_points_regular_grid( + points, + store / "points" / points_element / "points.parquet", + catalog=catalog, + grid=grid, + display_transform=transform, + feature_key=feature_key, + max_row_groups_per_file=max_row_groups_per_file, + compression=compression, + overwrite=True, + ) + # Paths in the manifest are relative to the profile directory, so Celldega can be + # pointed at that directory as its base_url with no reader change. + transcripts["directory"] = f"../../points/{points_element}/points.parquet" + + cell_segmentation = None + if shapes_element: + if shapes_element not in sdata.shapes: + raise ValueError(f"shapes element {shapes_element!r} not found; have {list(sdata.shapes)}") + shapes = sdata.shapes[shapes_element] + cell_segmentation = write_shapes_regular_grid( + shapes, + store / "shapes" / shapes_element / "shapes.parquet", + grid=grid, + display_transform=DisplayTransform.from_element(shapes, coordinate_system), + cell_index=list(table.obs_names) if table is not None else None, + max_row_groups_per_file=max_row_groups_per_file, + compression=compression, + overwrite=True, + ) + prefix = f"../../shapes/{shapes_element}" + if "path" in cell_segmentation: + cell_segmentation["path"] = f"{prefix}/{cell_segmentation['path']}" + else: + cell_segmentation["directory"] = f"{prefix}/{cell_segmentation['directory']}" + + cbg = None + if include_cbg and table is not None: + cbg = write_cbg_row_groups( + table, + profile_dir / "cbg", + catalog=catalog, + max_row_groups_per_file=max_row_groups_per_file, + compression=compression, + overwrite=True, + ) + + catalog.to_frame().to_parquet(profile_dir / "meta_gene.parquet", index=False) + + manifest = build_manifest( + grid=grid, + technology=technology, + transcripts=transcripts, + cell_segmentation=cell_segmentation, + cbg=cbg, + source={ + "store": store.name, + "points_element": points_element, + "shapes_element": shapes_element, + "table_element": table_element, + "coordinate_system": coordinate_system, + "tile_size_px": tile_size_px, + }, + ) + validate_manifest(manifest, base_path=profile_dir) + write_manifest(manifest, profile_dir) + return manifest + + +def xenium_spatially_tiled( + raw_path: str | Path, + output_path: str | Path, + *, + tile_size_px: float = 250.0, + max_row_groups_per_file: int = DEFAULT_MAX_ROW_GROUPS_PER_FILE, + include_cbg: bool = True, + compression: str = "zstd", + overwrite: bool = False, + **xenium_kwargs: Any, +) -> dict[str, Any]: + """Read raw Xenium data and write a spatially tiled SpatialData store in one call. + + Parameters + ---------- + raw_path + Directory of raw Xenium output. + output_path + Destination ``.zarr`` store. + tile_size_px + Tile edge length in display pixels. + max_row_groups_per_file + Row groups per chunk file. + include_cbg + Whether to write the gene-major cell-by-gene files. + compression + Parquet compression codec. + overwrite + Replace ``output_path`` if it exists. + xenium_kwargs + Forwarded to :func:`spatialdata_io.xenium`. + + Returns + ------- + The profile manifest. + """ + from spatialdata_io.readers.xenium import xenium + + output_path = Path(output_path) + if output_path.exists(): + if not overwrite: + raise FileExistsError(f"{output_path} exists; pass overwrite=True to replace it") + shutil.rmtree(output_path) + + sdata = xenium(raw_path, **xenium_kwargs) + sdata.write(output_path) + return add_spatial_tiling( + output_path, + tile_size_px=tile_size_px, + max_row_groups_per_file=max_row_groups_per_file, + include_cbg=include_cbg, + compression=compression, + ) diff --git a/tests/test_tiled_access.py b/tests/test_tiled_access.py new file mode 100644 index 00000000..cf61666b --- /dev/null +++ b/tests/test_tiled_access.py @@ -0,0 +1,240 @@ +"""Tests for the profile manifest and the opt-in tiling entry points.""" + +from __future__ import annotations + +import json +from pathlib import Path + +import geopandas as gpd +import numpy as np +import pandas as pd +import pyarrow.parquet as pq +import pytest +import scipy.sparse as sp +from anndata import AnnData +from shapely.geometry import Polygon + +from spatialdata_io.experimental.manifest import ( + MANIFEST_FILENAME, + PROFILE_NAME, + build_manifest, + validate_manifest, + write_manifest, +) +from spatialdata_io.experimental.regular_grid import RegularGrid +from spatialdata_io.experimental.tiled_access import add_spatial_tiling + +GRID = RegularGrid(origin_x=0.0, origin_y=0.0, tile_size_px=10.0, num_tiles_x=2, num_tiles_y=3) + + +# -- manifest ----------------------------------------------------------------- + + +def _transcripts_entry(**over): + entry = { + "directory": "trx", + "files": ["chunk_0.parquet"], + "max_row_groups_per_file": 400, + "total_row_groups": GRID.num_tiles, + "position_column": "display_xy", + "feature_column": "feature_code", + } + entry.update(over) + return entry + + +def test_manifest_uses_celldega_keys() -> None: + m = build_manifest(grid=GRID, transcripts=_transcripts_entry()) + for key in ("technology", "use_row_groups", "tile_grid", "row_group_files", "image_info"): + assert key in m + assert m["use_row_groups"] is True + assert m["profile"] == PROFILE_NAME + assert m["tile_grid"]["num_tiles_x"] == 2 + + +def test_manifest_omits_absent_elements() -> None: + m = build_manifest(grid=GRID, transcripts=_transcripts_entry()) + assert "cell_segmentation" not in m["row_group_files"] + assert "cbg" not in m["row_group_files"] + # images is always present so the reader's loop has something to iterate + assert m["row_group_files"]["images"] == {} + + +def test_validate_rejects_row_group_count_mismatch() -> None: + m = build_manifest(grid=GRID, transcripts=_transcripts_entry(total_row_groups=5)) + with pytest.raises(ValueError, match="does not match the tile grid"): + validate_manifest(m) + + +def test_validate_rejects_wrong_file_count() -> None: + m = build_manifest( + grid=GRID, transcripts=_transcripts_entry(max_row_groups_per_file=2, files=["chunk_0.parquet"]) + ) + with pytest.raises(ValueError, match="lists 1 file"): + validate_manifest(m) + + +def test_validate_rejects_entry_without_files_or_path() -> None: + m = build_manifest(grid=GRID, transcripts={"total_row_groups": 6, "max_row_groups_per_file": 400}) + with pytest.raises(ValueError, match="neither 'files' nor 'path'"): + validate_manifest(m) + + +def test_validate_rejects_empty_cbg_mapping() -> None: + m = build_manifest(grid=GRID, cbg={"directory": "cbg", "files": ["c.parquet"], "gene_to_row_group": {}}) + with pytest.raises(ValueError, match="gene_to_row_group is empty"): + validate_manifest(m) + + +def test_validate_rejects_cbg_row_group_out_of_range() -> None: + m = build_manifest( + grid=GRID, + cbg={ + "directory": "cbg", + "files": ["c.parquet"], + "max_row_groups_per_file": 2, + "total_row_groups": 3, + "gene_to_row_group": {"A": 0, "B": 1, "C": 9}, + }, + ) + with pytest.raises(ValueError, match="references row group 9"): + validate_manifest(m) + + +def test_validate_checks_files_exist(tmp_path: Path) -> None: + m = build_manifest(grid=GRID, transcripts=_transcripts_entry()) + with pytest.raises(ValueError, match="does not exist"): + validate_manifest(m, base_path=tmp_path) + + +def test_write_manifest_round_trips(tmp_path: Path) -> None: + m = build_manifest(grid=GRID, transcripts=_transcripts_entry()) + path = write_manifest(m, tmp_path) + assert path.name == MANIFEST_FILENAME + assert json.loads(path.read_text()) == m + + +# -- end-to-end tiling -------------------------------------------------------- + + +@pytest.fixture +def store(tmp_path: Path) -> Path: + """A minimal but complete SpatialData store: points, shapes and an annotating table.""" + from spatialdata import SpatialData + from spatialdata.models import PointsModel, ShapesModel, TableModel + from spatialdata.transformations import Scale, set_transformation + + rng = np.random.default_rng(0) + n = 200 + cells = [f"cell-{i}" for i in range(12)] + genes = ["GENEA", "GENEB", "GENEC"] + + pts = pd.DataFrame( + { + "x": rng.uniform(0, 9.9, n), + "y": rng.uniform(0, 14.9, n), + "feature_name": pd.Categorical(rng.choice([*genes, "NegControlProbe_0001"], n)), + "cell_id": rng.choice(cells, n), + } + ) + points = PointsModel.parse(pts, coordinates={"x": "x", "y": "y"}, feature_key="feature_name") + set_transformation(points, Scale([2.0, 2.0], axes=("x", "y")), "global") + + gdf = gpd.GeoDataFrame( + geometry=[ + Polygon([(x, y), (x + 0.5, y), (x + 0.5, y + 0.5), (x, y + 0.5)]) + for x, y in zip(rng.uniform(0.5, 9, 12), rng.uniform(0.5, 14, 12)) + ], + index=cells, + ) + shapes = ShapesModel.parse(gdf) + set_transformation(shapes, Scale([2.0, 2.0], axes=("x", "y")), "global") + + obs = pd.DataFrame({"region": pd.Categorical(["cell_boundaries"] * 12), "instance_id": range(12)}, index=cells) + table = TableModel.parse( + AnnData(X=sp.csr_matrix(rng.integers(0, 5, (12, 3)).astype(np.float32)), obs=obs, + var=pd.DataFrame(index=genes)), + region="cell_boundaries", region_key="region", instance_key="instance_id", + ) + + path = tmp_path / "s.zarr" + SpatialData(points={"transcripts": points}, shapes={"cell_boundaries": shapes}, + tables={"table": table}).write(path) + return path + + +def test_add_spatial_tiling_produces_a_valid_profile(store: Path) -> None: + manifest = add_spatial_tiling(store, tile_size_px=10.0) + profile = store / "visualization" / PROFILE_NAME + assert (profile / MANIFEST_FILENAME).exists() + assert (profile / "meta_gene.parquet").exists() + validate_manifest(manifest, base_path=profile) + + +def test_tiled_store_still_reads_with_read_zarr(store: Path) -> None: + import spatialdata + + before = spatialdata.read_zarr(store) + n_points, n_shapes = len(before.points["transcripts"]), len(before.shapes["cell_boundaries"]) + + add_spatial_tiling(store, tile_size_px=10.0) + + after = spatialdata.read_zarr(store) + assert len(after.points["transcripts"]) == n_points + assert len(after.shapes["cell_boundaries"]) == n_shapes + assert "display_xy" in after.points["transcripts"].columns + assert "display_geometry" in after.shapes["cell_boundaries"].columns + # canonical columns survive untouched + assert {"x", "y", "feature_name", "cell_id"} <= set(after.points["transcripts"].columns) + + +def test_manifest_paths_resolve_from_the_profile_directory(store: Path) -> None: + """Celldega is pointed at the profile dir, so relative paths must resolve from there.""" + manifest = add_spatial_tiling(store, tile_size_px=10.0) + profile = store / "visualization" / PROFILE_NAME + trx = manifest["row_group_files"]["transcripts"] + for f in trx["files"]: + assert (profile / trx["directory"] / f).resolve().exists() + + +def test_row_group_count_matches_grid(store: Path) -> None: + manifest = add_spatial_tiling(store, tile_size_px=10.0) + grid = RegularGrid.from_manifest_dict(manifest["tile_grid"]) + profile = store / "visualization" / PROFILE_NAME + trx = manifest["row_group_files"]["transcripts"] + total = sum(pq.ParquetFile(profile / trx["directory"] / f).metadata.num_row_groups for f in trx["files"]) + assert total == grid.num_tiles == manifest["row_group_files"]["transcripts"]["total_row_groups"] + + +def test_cbg_covers_every_gene(store: Path) -> None: + manifest = add_spatial_tiling(store, tile_size_px=10.0) + cbg = manifest["row_group_files"]["cbg"] + assert set(cbg["gene_to_row_group"]) == {"GENEA", "GENEB", "GENEC"} + assert "NegControlProbe_0001" not in cbg["gene_to_row_group"] + + +def test_tiling_is_rerunnable(store: Path) -> None: + """The two-call workflow means users will re-run this; it must not accumulate state.""" + first = add_spatial_tiling(store, tile_size_px=10.0) + second = add_spatial_tiling(store, tile_size_px=10.0) + assert first["tile_grid"] == second["tile_grid"] + assert first["row_group_files"]["transcripts"]["total_row_groups"] == ( + second["row_group_files"]["transcripts"]["total_row_groups"] + ) + + import spatialdata + + after = spatialdata.read_zarr(store) + cols = list(after.points["transcripts"].columns) + assert cols.count("display_xy") == 1 + assert cols.count("feature_code") == 1 + + +def test_missing_element_is_reported(store: Path) -> None: + with pytest.raises(ValueError, match="points element 'nope' not found"): + add_spatial_tiling(store, points_element="nope") + + +def test_cbg_can_be_skipped(store: Path) -> None: + manifest = add_spatial_tiling(store, tile_size_px=10.0, include_cbg=False) + assert "cbg" not in manifest["row_group_files"] From c62667631375d3eed8bdab1417542ee2488f71fb Mon Sep 17 00:00:00 2001 From: Nicolas Fernandez Date: Sun, 6 Sep 2026 15:30:20 -0400 Subject: [PATCH 07/17] docs: add the regular-grid tiled access protocol specification Normative description of celldega_regular_grid_v1: coordinate system, grid and tile numbering, row-group and multi-part file numbering, empty-tile handling, the render columns, feature and cell codes, CBG mapping, manifest layout, transport requirements, invalidation rules and a conformance checklist. Written viewer-independently with Celldega named as the reference client, and carrying the measurements behind the non-obvious choices (tile size, multi-file splitting, zstd) so the rationale survives the prototype. Also declares the 'columns' projection in the writer manifest fragments. Co-Authored-By: Claude Opus 5 (1M context) --- docs/design/regular_grid_tiled_access.md | 426 ++++++++++++++++++ .../experimental/points_parquet.py | 3 + .../experimental/shapes_parquet.py | 2 + 3 files changed, 431 insertions(+) create mode 100644 docs/design/regular_grid_tiled_access.md diff --git a/docs/design/regular_grid_tiled_access.md b/docs/design/regular_grid_tiled_access.md new file mode 100644 index 00000000..cb957eac --- /dev/null +++ b/docs/design/regular_grid_tiled_access.md @@ -0,0 +1,426 @@ +# Regular-grid tiled access profile (`celldega_regular_grid_v1`) + +**Status:** experimental · **Profile version:** 0.1.0 + +A viewer-independent protocol for fetching spatially-local subsets of a SpatialData store +over HTTP range requests, without downloading whole files and without consulting Parquet +statistics. + +Celldega is the first reference client. Nothing in this document is Celldega-specific; +another viewer (Vitessce, SpatialData.js, napari) could implement it from this text alone. + +--- + +## 1. Motivation and scope + +A viewer showing a 34,000 × 14,000 pixel tissue with 8 million transcripts must fetch only +what is on screen. Two things have to be true: + +1. The client can compute **which bytes it needs** from the viewport alone — no index + download, no metadata probing, no statistics. +2. Those bytes are **already in the form the GPU wants** — no per-point object + construction, no coordinate zipping, no WKB parsing. + +This profile achieves both by reordering rows into a deterministic grid of Parquet row +groups and adding a small number of render-oriented columns. + +**In scope:** transcript points, cell polygons, gene-major expression, image tiles, +and the manifest that describes them. + +**Out of scope:** clustering, linked views, annotation, neighbourhood analysis, and any +other viewer feature. This is a data-access protocol. + +### Design invariants + +- **The profile is opt-in.** Default SpatialData write behaviour is unchanged. +- **Canonical data is authoritative and preserved.** Render columns are additions; they + never replace canonical coordinates, identifiers, geometries or annotations. +- **A store carrying the profile is still an ordinary SpatialData store.** + `spatialdata.read_zarr()` works unchanged; a client that does not understand the profile + ignores the extra columns and the manifest. + +--- + +## 2. Coordinate system + +All profile coordinates are **level-0 pixels of a chosen reference image**, referred to +here as *display pixel space*. + +The mapping from canonical element coordinates to display pixels is the element's +SpatialData affine transformation into a named coordinate system (`global` by default). +It is recorded in the manifest: + +```json +"display_transform": { + "coordinate_space": "image-pixel", + "coordinate_system": "global", + "affine_matrix": [[4.70588235, 0.0, 0.0], [0.0, 4.70588235, 0.0]], + "rounding": "nearest" +} +``` + +Producers **MUST** apply the affine in float64 and round half-to-even (`rint`). Applying +it in float32 loses pixel accuracy on large images, and in some environments a float32 +array multiplied by a scalar stays float32. + +Display coordinates **MUST** be non-negative and **MUST** fit the declared integer dtype. +A producer encountering values outside that range **MUST** fail rather than clamp: an +out-of-range coordinate indicates a mismatched transform, not a rounding artefact. + +--- + +## 3. The grid + +A non-overlapping regular square grid over display pixel space, defined by five numbers: + +| field | meaning | +|---|---| +| `x_min`, `y_min` | grid origin, display pixels | +| `tile_size` | tile edge length, display pixels | +| `num_tiles_x`, `num_tiles_y` | grid dimensions | + +### Tile assignment + +``` +tile_x = floor((x_px - x_min) / tile_size) +tile_y = floor((y_px - y_min) / tile_size) +``` + +Tile bounds are **half-open** `[min, max)`: a coordinate lying exactly on an internal +boundary belongs to the **upper** tile. + +The single exception is the grid's outer edge. A coordinate equal to `x_max` or `y_max` +is clamped into the last tile, so that a point on the boundary of the dataset is not +lost. Coordinates beyond one tile past the extent are an error. + +### Tile numbering + +Tiles are numbered **x-major**: + +``` +tile_id = tile_x * num_tiles_y + tile_y +``` + +`tile_id` ranges over `[0, num_tiles_x * num_tiles_y)`. + +### Choosing `tile_size` + +`tile_size` is a tuning parameter, not a constant. It trades viewport granularity against +storage: smaller tiles fetch less off-screen data but fragment the file into more, +individually-compressed row groups. + +The recommended target is **roughly 20 cells per tile**, which is the granularity at +which a viewer fetches. On Xenium-density tissue this is about **250 display pixels**. +Measured on a Xenium human pancreas section (140,702 cells): + +| tile px | cells/tile | row groups | size vs untiled | +|---|---|---|---| +| 200 | 13.5 | 11,799 | +71% | +| **250** | **21.0** | **7,535** | **+60%** | +| 500 | 81.0 | 1,932 | +50% | + +--- + +## 4. Row groups and files + +### One tile, one row group + +Each logical tile is written as **exactly one Parquet row group**, at index `tile_id`. +Tiles containing no rows are written as **zero-row row groups**, not skipped. This is what +lets a client address a tile by formula with no lookup table. + +A conforming file therefore contains exactly `num_tiles_x * num_tiles_y` row groups across +all its parts. + +### Multi-part files + +Row groups are split across files: + +``` +file_index = tile_id // max_row_groups_per_file +local_row_group = tile_id % max_row_groups_per_file +``` + +`max_row_groups_per_file` defaults to **400**. + +Splitting is **required**, not cosmetic. A Parquet reader must fetch a file's entire +footer before reading any row group, and footer size grows with row-group count. On the +pancreas dataset a single 7,535-row-group file has a **7.4 MB footer**; split into 19 +files each footer is ~410 KB, and a client only fetches footers for files its viewport +actually touches. + +### File naming + +Chunk files are named `chunk_.parquet` with `` **zero-padded** to the width of the +largest index (`chunk_00.parquet` … `chunk_18.parquet`). + +Padding is required because consumers disagree about ordering: a client indexes the +manifest's `files` array by position, but tools that glob a directory sort +lexicographically, where `chunk_10` precedes `chunk_2`. Padding makes the two agree. + +> Existing Celldega DegaFiles use unpadded names. That is safe there because only the +> manifest-array consumer exists. Stores written under this profile use padded names. + +### Statistics + +Producers **SHOULD** write Parquet files with column statistics disabled. The tile formula +is the spatial index, so no conforming client reads column-chunk min/max, and statistics +inflate the footer the client must download before its first read. + +### Compression + +**zstd** is the recommended codec. On the pancreas dataset, snappy costs +38.5% over an +untiled store while zstd costs +4.9% for the same content, at no meaningful write cost. +Producers **MUST NOT** use a codec the target client cannot decode. + +--- + +## 5. Transcript points + +The canonical Points element gains two columns; every canonical column and the DataFrame +index are preserved. Physical row order changes (rows are grouped by tile), which is +permitted; rows **MUST NOT** be added, dropped or altered. + +### `display_xy` + +``` +fixed_size_list[2] +``` + +Integer display-pixel coordinates. The Arrow child buffer is therefore already +`[x0, y0, x1, y1, ...]`, directly usable as a deck.gl binary `getPosition` attribute. + +A client **MUST NOT** need to interleave separate x and y arrays. + +> A future revision may allow fixed-point sub-pixel coordinates +> (`stored = pixel * 16`, `scale = 0.0625`). Producers of v0.1.0 write integer pixels +> and declare `"scale": 1.0`. + +### `feature_code` + +``` +uint16 (or uint32 when the catalog exceeds 65535 entries) +``` + +An index into the feature catalog (§7). + +--- + +## 6. Cell polygons + +The canonical Shapes element gains two columns. The canonical geometry column and its +GeoParquet `geo` metadata are preserved, so the file remains readable by +`geopandas.read_parquet` and by SpatialData. + +### `display_geometry` + +``` +list[2]>> +``` + +Polygon → rings → interleaved integer pixel vertices. A client lifts `getPolygon` from the +flat coordinate child buffer and `startIndices` from the list offsets: + +``` +start_index[i] = ring_offsets[polygon_offsets[i]] +``` + +`display_geometry` is explicitly a **lossy display representation**: it holds the exterior +ring only, and for a MultiPolygon only the largest part. The canonical geometry is retained +alongside it and is authoritative. + +> **Interoperability warning.** GeoArrow permits `struct` coordinates as well as +> interleaved `fixed_size_list`. A client walking the nesting blindly will, on struct +> coordinates, obtain the `x` child alone and render wrong polygons with no error. Clients +> **MUST** verify the vertex level is a `FixedSizeList` before treating the buffer as +> interleaved. (`geopandas.to_parquet` emits struct coordinates, so this is reachable in +> practice.) + +### `cell_code` + +``` +uint32 +``` + +Positional index into the annotating table's `obs` order, so cells can be coloured from an +expression vector without a string join in the client. + +### Tile assignment + +A cell is assigned to **exactly one tile**, by its **centroid** in display pixel space. + +A polygon whose outline crosses into neighbouring tiles is **NOT** duplicated into them. +Duplication would inflate the file and make cell counts wrong. A client rendering a +viewport should expect polygons to overhang tile boundaries and, if it needs full coverage +at the edges, fetch one extra ring of tiles. + +--- + +## 7. Feature catalog + +An ordered vocabulary mapping feature names to `feature_code`, stored as +`meta_gene.parquet` with columns `name`, `feature_code`, `is_gene`. + +Ordering is normative: + +1. Codes `[0, n_genes)` are **genes**, in the annotating table's `var_names` order. +2. Codes `>= n_genes` are **non-gene features** (negative controls, unassigned codewords), + sorted for reproducibility. + +Because genes come first and in table order, **a gene's `feature_code` is also its CBG row +group** (§8). One integer addresses both a transcript's identity and its expression vector. + +Non-gene features **MUST** be retained and **MUST NOT** be folded into a gene; doing so +would fabricate expression. `n_genes` is recorded in the manifest so a client can tell the +two apart. + +--- + +## 8. Cell-by-gene expression + +Gene-major, one row group per gene, so selecting a gene fetches one row group and touches +no transcript data. + +Columns: + +| column | type | meaning | +|---|---|---| +| `cell_id` | uint32 | cell code (§6), not a string barcode | +| `expression` | float32 | non-zero value | +| `gene` | string | gene name | + +Zero values are omitted, including sparse *stored* zeros. + +Row group `i` holds catalog gene `i`, including genes with no expression (written as a +zero-row row group), preserving the `feature_code == row group` invariant. The explicit +mapping is written both in the manifest and in the Parquet schema metadata: + +``` +gene_to_row_group JSON object +num_genes integer +storage_mode "row_groups_cbg_chunked" +``` + +A client **MAY** use the mapping rather than assuming the identity. + +> This is a **transpose**, not a redundant copy. SpatialData tables are stored cell-major +> (CSR); assembling one gene's vector from them requires reading the whole matrix or doing +> one random access per cell. + +--- + +## 9. Image tiles + +*(Not implemented in v0.1.0; specified here for forward compatibility.)* + +Canonical OME-Zarr images are retained and authoritative. An optional derived WebP pyramid +may be stored as Parquet row groups with columns `zoom`, `tile_x`, `tile_y`, `image_data` +(encoded WebP bytes). + +Image tiles at zoom 0 **MUST** use the same level-0 pixel coordinate system as +`display_xy` and `display_geometry`. + +Per channel the manifest records: source image element, source dimensions and dtype, +reference pyramid level, tile size, per-zoom grid dimensions, display intensity min/max, +gamma, colour, downsampling method, and WebP lossless/lossy setting. + +--- + +## 10. The manifest + +A JSON document named `landscape_parameters.json`, conventionally at +`.zarr/visualization/celldega_regular_grid_v1/`. + +Paths inside it are **relative to the manifest's own directory**, so a client pointed at +that directory resolves into the store without knowing the zarr layout: + +```json +{ + "technology": "Xenium", + "use_row_groups": true, + "profile": "celldega_regular_grid_v1", + "profile_version": "0.1.0", + "tile_grid": { + "num_tiles_x": 137, "num_tiles_y": 55, "tile_size": 250.0, + "x_min": 0.0, "y_min": 0.0, "x_max": 34250.0, "y_max": 13750.0 + }, + "row_group_files": { + "transcripts": { + "directory": "../../points/transcripts/points.parquet", + "files": ["chunk_00.parquet", "..."], + "max_row_groups_per_file": 400, + "total_row_groups": 7535, + "position_column": "display_xy", + "feature_column": "feature_code", + "columns": ["display_xy", "feature_code"] + }, + "cell_segmentation": { + "directory": "../../shapes/cell_boundaries/shapes.parquet", + "geometry_column": "display_geometry", + "cell_id_column": "cell_code", + "columns": ["display_geometry", "cell_code"] + }, + "cbg": { "directory": "cbg", "gene_to_row_group": {} }, + "images": {} + }, + "image_info": [] +} +``` + +`columns` is the projection a client should request. Honouring it is what keeps canonical +coordinates, identifiers and QC columns off the wire during rendering. + +A client that does not recognise the profile-specific keys and finds no column names +declared **SHOULD** fall back to reading all columns. + +--- + +## 11. Transport requirements + +A conforming host **MUST** support: + +- **HTTP range requests**, answering `Range: bytes=a-b` with `206 Partial Content` and a + correct `Content-Range`. +- **Suffix ranges** (`Range: bytes=-8`), which is how a reader locates the Parquet footer. +- **CORS**, with `Access-Control-Allow-Origin` and `Access-Control-Expose-Headers` + including `Content-Range`, for browser clients on another origin. + +Hugging Face dataset `resolve/` URLs satisfy all three. + +--- + +## 12. Invalidation + +The profile is derived data and **MUST** be regenerated when any of the following change: + +- transcript rows are added, removed, or spatially moved; +- the feature catalog or `var_names` order changes; +- canonical cell identifiers or table row order change (invalidates `cell_code`); +- cell geometries change; +- the reference image, or the transform into display pixel space, changes; +- `tile_size`, the grid origin, or `max_row_groups_per_file` change; +- image intensity windowing or WebP settings change (images only). + +Changing a SpatialData transformation that does **not** affect the declared display +coordinate system does not require regeneration, but this **MUST** be verified rather than +assumed — compare the resulting affine against `display_transform.affine_matrix`. + +Producers **SHOULD** record a source fingerprint in `source` so staleness is detectable. + +--- + +## 13. Conformance checklist + +A producer conforms if: + +- [ ] total row groups equals `num_tiles_x * num_tiles_y`, empty tiles included +- [ ] row group index equals `tile_x * num_tiles_y + tile_y` +- [ ] every canonical row appears exactly once; index preserved +- [ ] canonical columns and geometries are bit-identical to the source +- [ ] `display_xy` is `fixed_size_list[2]` with an interleaved child buffer +- [ ] `display_geometry` vertices are `fixed_size_list`, not `struct` +- [ ] each cell appears exactly once, in its centroid's tile +- [ ] `feature_code` matches the catalog; genes precede non-genes +- [ ] CBG row group index equals `feature_code` for every gene +- [ ] the store still opens with `spatialdata.read_zarr()` +- [ ] the manifest validates and every declared file exists diff --git a/src/spatialdata_io/experimental/points_parquet.py b/src/spatialdata_io/experimental/points_parquet.py index 4bb11d2a..eea0daa0 100644 --- a/src/spatialdata_io/experimental/points_parquet.py +++ b/src/spatialdata_io/experimental/points_parquet.py @@ -285,6 +285,9 @@ def write_points_regular_grid( "position_dtype": "uint32", "position_size": 2, "feature_column": FEATURE_COLUMN, + # Projected by the client, so canonical coordinates, ids and QC columns are never + # decoded or transferred during ordinary rendering. + "columns": [POSITION_COLUMN, FEATURE_COLUMN], "n_rows": int(table.num_rows), "tile_grid": grid.to_manifest_dict(), "display_transform": transform.to_manifest_dict(), diff --git a/src/spatialdata_io/experimental/shapes_parquet.py b/src/spatialdata_io/experimental/shapes_parquet.py index b5a2122e..8054c604 100644 --- a/src/spatialdata_io/experimental/shapes_parquet.py +++ b/src/spatialdata_io/experimental/shapes_parquet.py @@ -241,6 +241,8 @@ def write_shapes_regular_grid( fragment: dict[str, Any] = { "geometry_column": GEOMETRY_COLUMN, "cell_id_column": CELL_CODE_COLUMN, + # Projected by the client: the canonical WKB geometry is never transferred. + "columns": [GEOMETRY_COLUMN, CELL_CODE_COLUMN], "max_row_groups_per_file": max_row_groups_per_file, "total_row_groups": grid.num_tiles, "n_shapes": int(table.num_rows), From 34c52eace5441dd8d43d99dc0367522de4a51ab1 Mon Sep 17 00:00:00 2001 From: Nicolas Fernandez Date: Sun, 6 Sep 2026 16:06:38 -0400 Subject: [PATCH 08/17] feat(experimental): add WebP display pyramid and wire images into the profile Writes a browser-ready WebP pyramid as Parquet row groups in Celldega's existing ImageRowGroupReader layout: zoom/tile_x/tile_y/image_data, one tile per row group, zoom_info in the schema metadata. Zoom numbering follows DeepZoom and was verified tile-for-tile against 'vips dzsave' on the same input: levels 0..11 for a 1300x700 image, level 11 = 3x2, level 10 = 2x1. So the output is interchangeable with tiles Celldega already produces. On Xenium pancreas it independently derives max_pyramid_zoom=16 and image_dimensions 34155x13770, matching the DegaFiles reference exactly. Encoding uses Pillow rather than libvips, so spatialdata-io gains no system dependency. Two memory measures matter on real images: the display window is taken from the smallest multiscale level rather than by scanning a 500-megapixel plane, and the 8-bit conversion runs in row blocks instead of materializing a float copy of the whole image. The canonical OME-Zarr image is untouched; this is explicitly a display cache, and the source element, dimensions, dtype, window, gamma, downsampling method and WebP settings are all recorded for invalidation. Also aligns the manifest with a real DegaFiles landscape_parameters.json (max_pyramid_zoom, image_dimensions, image_format, use_int_index, segmentation_approach, top-level tile_size). Co-Authored-By: Claude Opus 5 (1M context) --- .../experimental/feature_catalog.py | 7 +- src/spatialdata_io/experimental/manifest.py | 12 + .../experimental/points_parquet.py | 1 + .../experimental/regular_grid.py | 4 +- .../experimental/tiled_access.py | 55 ++++ .../experimental/webp_parquet.py | 310 ++++++++++++++++++ tests/test_points_parquet.py | 4 +- tests/test_shapes_parquet.py | 4 +- tests/test_tiled_access.py | 21 +- tests/test_webp_parquet.py | 198 +++++++++++ 10 files changed, 595 insertions(+), 21 deletions(-) create mode 100644 src/spatialdata_io/experimental/webp_parquet.py create mode 100644 tests/test_webp_parquet.py diff --git a/src/spatialdata_io/experimental/feature_catalog.py b/src/spatialdata_io/experimental/feature_catalog.py index 09ed7d29..1b2664de 100644 --- a/src/spatialdata_io/experimental/feature_catalog.py +++ b/src/spatialdata_io/experimental/feature_catalog.py @@ -123,9 +123,10 @@ def from_points_and_table( if hasattr(col, "cat"): try: features = list(col.cat.categories) - except Exception: - # Unknown categories (typical straight after read_zarr): realize just the - # category list, which is tiny, rather than computing the whole column. + except (NotImplementedError, AttributeError): + # dask raises AttributeNotImplementedError (a subclass of both) for + # unknown categories, which is the normal state straight after read_zarr. + # Realize just the category list, which is tiny, rather than the column. features = list(col.cat.as_known().cat.categories) else: features = list(col.unique().compute() if hasattr(col, "compute") else col.unique()) diff --git a/src/spatialdata_io/experimental/manifest.py b/src/spatialdata_io/experimental/manifest.py index 58f4c9af..778ebcd3 100644 --- a/src/spatialdata_io/experimental/manifest.py +++ b/src/spatialdata_io/experimental/manifest.py @@ -43,6 +43,8 @@ def build_manifest( cbg: dict[str, Any] | None = None, images: dict[str, Any] | None = None, image_info: list[dict[str, Any]] | None = None, + image_dimensions: dict[str, Any] | None = None, + max_pyramid_zoom: int | None = None, source: dict[str, Any] | None = None, ) -> dict[str, Any]: """Assemble the profile manifest from the fragments returned by each writer. @@ -77,15 +79,25 @@ def build_manifest( manifest: dict[str, Any] = { # -- keys Celldega's existing reader consumes ------------------------- + # Names and shapes follow a DegaFiles landscape_parameters.json so the reader + # needs no special-casing for a SpatialData store. "technology": technology, "use_row_groups": True, + "use_int_index": True, + "segmentation_approach": ["default"], + "tile_size": grid.tile_size_px, "tile_grid": grid.to_manifest_dict(), "row_group_files": row_group_files, "image_info": image_info or [], + "image_format": ".webp", # -- profile identification ------------------------------------------- "profile": PROFILE_NAME, "profile_version": PROFILE_VERSION, } + if image_dimensions is not None: + manifest["image_dimensions"] = image_dimensions + if max_pyramid_zoom is not None: + manifest["max_pyramid_zoom"] = max_pyramid_zoom if source is not None: manifest["source"] = source return manifest diff --git a/src/spatialdata_io/experimental/points_parquet.py b/src/spatialdata_io/experimental/points_parquet.py index eea0daa0..b68ce458 100644 --- a/src/spatialdata_io/experimental/points_parquet.py +++ b/src/spatialdata_io/experimental/points_parquet.py @@ -86,6 +86,7 @@ def apply(self, x: NDArray[Any], y: NDArray[Any]) -> tuple[NDArray[np.float64], return a * x + b * y + c, d * x + e * y + f def to_manifest_dict(self) -> dict[str, Any]: + """Serialize the transform for the manifest, so a client can reproduce the mapping.""" return { "coordinate_space": "image-pixel", "coordinate_system": self.coordinate_system, diff --git a/src/spatialdata_io/experimental/regular_grid.py b/src/spatialdata_io/experimental/regular_grid.py index 84a89b76..4d19334d 100644 --- a/src/spatialdata_io/experimental/regular_grid.py +++ b/src/spatialdata_io/experimental/regular_grid.py @@ -18,8 +18,8 @@ exactly one Parquet row group, including empty tiles, which are written as zero-row row groups. Row groups are then split across files:: - file_index = tile_id // max_row_groups_per_file - local_row_group = tile_id % max_row_groups_per_file + file_index = tile_id // max_row_groups_per_file + local_row_group = tile_id % max_row_groups_per_file Multi-file output is deliberate: a Parquet reader must fetch a file's entire footer before it can read any row group, and footer size grows with row-group count. Splitting diff --git a/src/spatialdata_io/experimental/tiled_access.py b/src/spatialdata_io/experimental/tiled_access.py index 53fb1fd8..b696278e 100644 --- a/src/spatialdata_io/experimental/tiled_access.py +++ b/src/spatialdata_io/experimental/tiled_access.py @@ -78,6 +78,12 @@ def add_spatial_tiling( feature_key: str = "feature_name", technology: str = "Xenium", include_cbg: bool = True, + image_element: str | None = None, + image_channel: int | str | None = None, + image_name: str = "dapi", + image_button_name: str = "DAPI", + image_color: tuple[int, int, int] = (0, 0, 255), + image_tile_size: int = 512, compression: str = "zstd", ) -> dict[str, Any]: """Add the regular-grid visualization profile to an existing SpatialData store. @@ -109,6 +115,15 @@ def add_spatial_tiling( Celldega technology string recorded in the manifest. include_cbg Whether to write the gene-major cell-by-gene files. + image_element + Name of an image element to render into a WebP display pyramid, or ``None`` to + skip images. The canonical OME-Zarr image is left untouched either way. + image_channel + Channel index or name to render. Defaults to the first channel. + image_name, image_button_name, image_color + Celldega channel descriptor recorded in ``image_info``. + image_tile_size + Image tile edge length in pixels. compression Parquet compression codec. @@ -186,6 +201,35 @@ def add_spatial_tiling( overwrite=True, ) + images: dict[str, Any] = {} + image_info: list[dict[str, Any]] = [] + image_dimensions: dict[str, Any] | None = None + max_pyramid_zoom: int | None = None + if image_element: + from spatialdata_io.experimental.webp_parquet import write_webp_pyramid + + if image_element not in sdata.images: + raise ValueError(f"image element {image_element!r} not found; have {list(sdata.images)}") + pyramid = write_webp_pyramid( + sdata.images[image_element], + profile_dir / "images" / image_name, + channel=image_channel, + tile_size=image_tile_size, + source_element=image_element, + overwrite=True, + ) + # ImageRowGroupReader resolves files as baseUrl/directory/file and reads the + # per-zoom grid from the entry's zoom_info. + pyramid["directory"] = f"images/{image_name}" + images[image_name] = pyramid + image_info = [{"name": image_name, "button_name": image_button_name, "color": list(image_color)}] + image_dimensions = { + "width": pyramid["source_width"], + "height": pyramid["source_height"], + "tile_size": image_tile_size, + } + max_pyramid_zoom = pyramid["max_zoom"] + catalog.to_frame().to_parquet(profile_dir / "meta_gene.parquet", index=False) manifest = build_manifest( @@ -194,11 +238,16 @@ def add_spatial_tiling( transcripts=transcripts, cell_segmentation=cell_segmentation, cbg=cbg, + images=images, + image_info=image_info, + image_dimensions=image_dimensions, + max_pyramid_zoom=max_pyramid_zoom, source={ "store": store.name, "points_element": points_element, "shapes_element": shapes_element, "table_element": table_element, + "image_element": image_element, "coordinate_system": coordinate_system, "tile_size_px": tile_size_px, }, @@ -217,6 +266,7 @@ def xenium_spatially_tiled( include_cbg: bool = True, compression: str = "zstd", overwrite: bool = False, + tiling: dict[str, Any] | None = None, **xenium_kwargs: Any, ) -> dict[str, Any]: """Read raw Xenium data and write a spatially tiled SpatialData store in one call. @@ -237,6 +287,10 @@ def xenium_spatially_tiled( Parquet compression codec. overwrite Replace ``output_path`` if it exists. + tiling + Extra keyword arguments for :func:`add_spatial_tiling`, for example + ``{"image_element": "morphology_focus", "image_channel": "DAPI"}``. Kept separate + from ``xenium_kwargs`` because the two functions have distinct option sets. xenium_kwargs Forwarded to :func:`spatialdata_io.xenium`. @@ -260,4 +314,5 @@ def xenium_spatially_tiled( max_row_groups_per_file=max_row_groups_per_file, include_cbg=include_cbg, compression=compression, + **(tiling or {}), ) diff --git a/src/spatialdata_io/experimental/webp_parquet.py b/src/spatialdata_io/experimental/webp_parquet.py new file mode 100644 index 00000000..5c39c5b5 --- /dev/null +++ b/src/spatialdata_io/experimental/webp_parquet.py @@ -0,0 +1,310 @@ +"""Write a browser-ready WebP image pyramid as Parquet row groups. + +The canonical OME-Zarr image stays authoritative; this is a derived *display cache*, and +must never be treated as the quantitative image. + +Layout matches Celldega's existing ``ImageRowGroupReader``: columns ``zoom``, ``tile_x``, +``tile_y``, ``image_data`` (encoded WebP bytes), one tile per row group, with a +``zoom_info`` map in the Parquet schema metadata:: + + row_group_index = zoom_info[zoom].row_group_offset + tile_x * num_tiles_y + tile_y + +Zoom numbering follows the DeepZoom convention, so the same numbers work against tiles +produced by ``vips dzsave``: level ``max_zoom`` is full resolution and each level below +halves both dimensions, down to a level that fits in a single tile. + +Encoding uses Pillow rather than libvips, so this adds no system dependency. Pyramid +levels are produced from the SpatialData image's own multiscale levels where they line up, +and by downsampling otherwise. +""" + +from __future__ import annotations + +import io +import json +import math +import shutil +from pathlib import Path +from typing import Any + +import numpy as np +import pyarrow as pa +import pyarrow.parquet as pq +from numpy.typing import NDArray + +__all__ = ["write_webp_pyramid", "DEFAULT_IMAGE_TILE_SIZE"] + +#: DeepZoom tile size used by Celldega's image pipeline. +DEFAULT_IMAGE_TILE_SIZE = 512 + +#: Image tiles are numerous and individually small, so more fit comfortably per file +#: than for transcripts. +DEFAULT_IMAGE_ROW_GROUPS_PER_FILE = 2000 + + +def _require_pillow() -> Any: + try: + from PIL import Image, features + except ImportError as exc: + raise RuntimeError("writing a WebP pyramid requires Pillow: pip install 'Pillow>=10'") from exc + if not features.check("webp"): + raise RuntimeError("this Pillow build has no WebP support; reinstall Pillow with WebP enabled") + return Image + + +def _select_channel(array: Any, channel: int | str | None) -> Any: + """Reduce one multiscale level to a 2D (y, x) plane.""" + if hasattr(array, "data_vars"): + array = array[next(iter(array.data_vars))] + dims = getattr(array, "dims", None) + if dims and "c" in dims: + if channel is None: + array = array.isel(c=0) + elif isinstance(channel, str): + array = array.sel(c=channel) + else: + array = array.isel(c=channel) + return array + + +def _as_2d(image: Any, channel: int | str | None) -> tuple[Any, Any]: + """Return ``(full_resolution_plane, coarse_plane)`` for a SpatialData image element. + + The coarse plane is the smallest available multiscale level, used only to choose the + display intensity window cheaply. For a single-scale image both are the same array. + """ + if hasattr(image, "children") and len(image.children): + levels = list(image.children) + return ( + _select_channel(image[levels[0]], channel), + _select_channel(image[levels[-1]], channel), + ) + plane = _select_channel(image, channel) + return plane, plane + + +def _display_window(sample_source: Any, display_min: float | None, display_max: float | None) -> tuple[float, float]: + """Choose the intensity window from a (possibly coarse) sample of the image. + + The window is a display choice, so it is taken from a downsampled level rather than by + scanning a 500-megapixel full-resolution plane. + """ + if display_min is not None and display_max is not None: + lo, hi = float(display_min), float(display_max) + else: + sample = np.asarray(sample_source) + sample = sample[np.isfinite(sample)] if sample.dtype.kind == "f" else sample.ravel() + lo = float(display_min) if display_min is not None else float(np.percentile(sample, 1.0)) + hi = float(display_max) if display_max is not None else float(np.percentile(sample, 99.9)) + return (lo, hi if hi > lo else lo + 1.0) + + +def _to_uint8(plane: NDArray[Any], lo: float, hi: float, gamma: float, block_rows: int = 2048) -> NDArray[np.uint8]: + """Window a plane to 8-bit in row blocks, avoiding a float copy of the whole image.""" + height = plane.shape[0] + out = np.empty(plane.shape, dtype=np.uint8) + inv_gamma = 1.0 / gamma + for start in range(0, height, block_rows): + stop = min(start + block_rows, height) + block = np.asarray(plane[start:stop], dtype=np.float32) + block -= lo + block /= hi - lo + np.clip(block, 0.0, 1.0, out=block) + if gamma != 1.0: + block **= inv_gamma + block *= 255.0 + block += 0.5 + out[start:stop] = block.astype(np.uint8) + return out + + +def _downsample_half(a: NDArray[np.uint8]) -> NDArray[np.uint8]: + """Box-filter by 2 in both axes, padding odd edges by replication.""" + h, w = a.shape + if h % 2: + a = np.vstack([a, a[-1:]]) + if w % 2: + a = np.hstack([a, a[:, -1:]]) + return a.reshape(a.shape[0] // 2, 2, a.shape[1] // 2, 2).mean(axis=(1, 3)).astype(np.uint8) + + +def write_webp_pyramid( + image: Any, + output_dir: str | Path, + *, + channel: int | str | None = None, + tile_size: int = DEFAULT_IMAGE_TILE_SIZE, + display_min: float | None = None, + display_max: float | None = None, + gamma: float = 1.0, + quality: int = 85, + lossless: bool = False, + max_row_groups_per_file: int = DEFAULT_IMAGE_ROW_GROUPS_PER_FILE, + source_element: str = "", + overwrite: bool = False, +) -> dict[str, Any]: + """Write a DeepZoom-numbered WebP pyramid as Parquet row groups. + + Parameters + ---------- + image + A SpatialData image element (``DataTree`` or ``DataArray``). + output_dir + Directory to write chunk files into. Written atomically. + channel + Channel index or name to render. Defaults to the first channel. + tile_size + Tile edge length in pixels. + display_min, display_max + Intensity window. Defaults to the 1st and 99.9th percentiles of the full-resolution + plane, which is a display choice and is recorded in the returned metadata. + gamma + Display gamma applied after windowing. + quality + WebP quality when ``lossless`` is False. + lossless + Whether to encode losslessly. + max_row_groups_per_file + Tiles per chunk file. + source_element + Name of the canonical image element, recorded for invalidation. + overwrite + Replace ``output_dir`` if it exists. + + Returns + ------- + The manifest fragment describing the written pyramid. + """ + Image = _require_pillow() + + output_dir = Path(output_dir) + if output_dir.exists() and not overwrite: + raise FileExistsError(f"{output_dir} exists; pass overwrite=True to replace it") + + plane, coarse = _as_2d(image, channel) + if plane.ndim != 2: + raise ValueError(f"expected a 2D plane after channel selection, got shape {plane.shape}") + source_dtype = str(plane.dtype) + + applied_min, applied_max = _display_window(coarse, display_min, display_max) + full = _to_uint8(plane, applied_min, applied_max, gamma) + height, width = full.shape + + # DeepZoom numbering: level max_zoom is full resolution and each level below halves + # both dimensions, down to level 0 (a single pixel). Every level is generated so the + # numbering matches `vips dzsave` output exactly. + max_zoom = max(1, math.ceil(math.log2(max(width, height)))) + levels: dict[int, NDArray[np.uint8]] = {max_zoom: full} + current = full + for zoom in range(max_zoom - 1, -1, -1): + current = _downsample_half(current) + levels[zoom] = current + + schema = pa.schema( + [ + pa.field("zoom", pa.int32()), + pa.field("tile_x", pa.int32()), + pa.field("tile_y", pa.int32()), + pa.field("image_data", pa.binary()), + ] + ) + + # Enumerate tiles in the reader's order: zoom ascending, then column-major within zoom. + ordered: list[tuple[int, int, int]] = [] + zoom_info: dict[str, dict[str, int]] = {} + for zoom in sorted(levels): + lh, lw = levels[zoom].shape + nx = max(1, math.ceil(lw / tile_size)) + ny = max(1, math.ceil(lh / tile_size)) + zoom_info[str(zoom)] = { + "num_tiles_x": nx, + "num_tiles_y": ny, + "num_tiles": nx * ny, + "row_group_offset": len(ordered), + } + ordered.extend((zoom, tx, ty) for tx in range(nx) for ty in range(ny)) + + n_files = max(1, -(-len(ordered) // max_row_groups_per_file)) + width_digits = len(str(n_files - 1)) if n_files > 1 else 1 + filenames = [f"chunk_{i:0{width_digits}d}.parquet" for i in range(n_files)] + + schema = schema.with_metadata( + { + b"zoom_info": json.dumps(zoom_info).encode(), + b"storage_mode": b"row_groups_image_chunked", + b"max_row_groups_per_file": str(max_row_groups_per_file).encode(), + b"tile_size": str(tile_size).encode(), + b"profile": b"celldega_regular_grid_v1", + } + ) + + staging = output_dir.with_name(output_dir.name + ".tmp") + if staging.exists(): + shutil.rmtree(staging) + staging.mkdir(parents=True) + + encode_kwargs: dict[str, Any] = {"format": "WEBP", "lossless": lossless} + if not lossless: + encode_kwargs["quality"] = quality + + try: + writer: pq.ParquetWriter | None = None + current_file = -1 + for index, (zoom, tx, ty) in enumerate(ordered): + file_index = index // max_row_groups_per_file + if file_index != current_file: + if writer is not None: + writer.close() + writer = pq.ParquetWriter(staging / filenames[file_index], schema, write_statistics=False) + current_file = file_index + + level = levels[zoom] + crop = level[ty * tile_size : (ty + 1) * tile_size, tx * tile_size : (tx + 1) * tile_size] + buf = io.BytesIO() + Image.fromarray(crop, mode="L").save(buf, **encode_kwargs) + + assert writer is not None + writer.write_table( + pa.table( + { + "zoom": pa.array([zoom], pa.int32()), + "tile_x": pa.array([tx], pa.int32()), + "tile_y": pa.array([ty], pa.int32()), + "image_data": pa.array([buf.getvalue()], pa.binary()), + }, + schema=schema, + ) + ) + if writer is not None: + writer.close() + except BaseException: + shutil.rmtree(staging, ignore_errors=True) + raise + + if output_dir.exists(): + shutil.rmtree(output_dir) + staging.rename(output_dir) + + return { + "directory": output_dir.name, + "files": filenames, + "max_row_groups_per_file": max_row_groups_per_file, + "total_row_groups": len(ordered), + "zoom_info": zoom_info, + "min_zoom": min(levels), + "max_zoom": max_zoom, + "tile_size": tile_size, + "image_format": ".webp", + # Recorded so the display cache can be invalidated when any of it changes. + "source_element": source_element, + "source_width": int(width), + "source_height": int(height), + "source_dtype": source_dtype, + "channel": channel, + "display_min": applied_min, + "display_max": applied_max, + "gamma": gamma, + "downsampling": "box-2x2-mean", + "webp_lossless": lossless, + "webp_quality": None if lossless else quality, + } diff --git a/tests/test_points_parquet.py b/tests/test_points_parquet.py index 18d93dd2..d8105cbf 100644 --- a/tests/test_points_parquet.py +++ b/tests/test_points_parquet.py @@ -141,9 +141,7 @@ def test_canonical_columns_are_unchanged(written: tuple[Path, dict], points: pd. merged = got.set_index("transcript_id").loc[original["transcript_id"].to_numpy()] for col in ("x", "y", "z", "cell_id", "qv"): - np.testing.assert_array_equal( - merged[col].to_numpy(), original[col].to_numpy(), err_msg=f"column {col} changed" - ) + np.testing.assert_array_equal(merged[col].to_numpy(), original[col].to_numpy(), err_msg=f"column {col} changed") assert list(merged["feature_name"].astype(str)) == list(original["feature_name"].astype(str)) diff --git a/tests/test_shapes_parquet.py b/tests/test_shapes_parquet.py index 16bf68e5..e30a8e20 100644 --- a/tests/test_shapes_parquet.py +++ b/tests/test_shapes_parquet.py @@ -31,9 +31,7 @@ def _square(cx: float, cy: float, half: float) -> Polygon: """A square in *canonical* coords, centred on (cx, cy).""" - return Polygon( - [(cx - half, cy - half), (cx + half, cy - half), (cx + half, cy + half), (cx - half, cy + half)] - ) + return Polygon([(cx - half, cy - half), (cx + half, cy - half), (cx + half, cy + half), (cx - half, cy + half)]) #: name -> (geometry, expected tile id). Canonical coords; pixels are 2x. diff --git a/tests/test_tiled_access.py b/tests/test_tiled_access.py index cf61666b..a1f8eb23 100644 --- a/tests/test_tiled_access.py +++ b/tests/test_tiled_access.py @@ -67,9 +67,7 @@ def test_validate_rejects_row_group_count_mismatch() -> None: def test_validate_rejects_wrong_file_count() -> None: - m = build_manifest( - grid=GRID, transcripts=_transcripts_entry(max_row_groups_per_file=2, files=["chunk_0.parquet"]) - ) + m = build_manifest(grid=GRID, transcripts=_transcripts_entry(max_row_groups_per_file=2, files=["chunk_0.parquet"])) with pytest.raises(ValueError, match="lists 1 file"): validate_manifest(m) @@ -152,14 +150,16 @@ def store(tmp_path: Path) -> Path: obs = pd.DataFrame({"region": pd.Categorical(["cell_boundaries"] * 12), "instance_id": range(12)}, index=cells) table = TableModel.parse( - AnnData(X=sp.csr_matrix(rng.integers(0, 5, (12, 3)).astype(np.float32)), obs=obs, - var=pd.DataFrame(index=genes)), - region="cell_boundaries", region_key="region", instance_key="instance_id", + AnnData( + X=sp.csr_matrix(rng.integers(0, 5, (12, 3)).astype(np.float32)), obs=obs, var=pd.DataFrame(index=genes) + ), + region="cell_boundaries", + region_key="region", + instance_key="instance_id", ) path = tmp_path / "s.zarr" - SpatialData(points={"transcripts": points}, shapes={"cell_boundaries": shapes}, - tables={"table": table}).write(path) + SpatialData(points={"transcripts": points}, shapes={"cell_boundaries": shapes}, tables={"table": table}).write(path) return path @@ -218,8 +218,9 @@ def test_tiling_is_rerunnable(store: Path) -> None: first = add_spatial_tiling(store, tile_size_px=10.0) second = add_spatial_tiling(store, tile_size_px=10.0) assert first["tile_grid"] == second["tile_grid"] - assert first["row_group_files"]["transcripts"]["total_row_groups"] == ( - second["row_group_files"]["transcripts"]["total_row_groups"] + assert ( + first["row_group_files"]["transcripts"]["total_row_groups"] + == (second["row_group_files"]["transcripts"]["total_row_groups"]) ) import spatialdata diff --git a/tests/test_webp_parquet.py b/tests/test_webp_parquet.py new file mode 100644 index 00000000..2198fc26 --- /dev/null +++ b/tests/test_webp_parquet.py @@ -0,0 +1,198 @@ +"""Tests for the WebP display pyramid. + +Zoom numbering follows DeepZoom so the output is interchangeable with `vips dzsave` +tiles; the numbering tests below encode that convention explicitly. +""" + +from __future__ import annotations + +import io +import json +import math +from pathlib import Path + +import numpy as np +import pyarrow.parquet as pq +import pytest +import xarray as xr + +from spatialdata_io.experimental.webp_parquet import write_webp_pyramid + +TILE = 128 +WIDTH, HEIGHT = 500, 300 + + +def _image(width: int = WIDTH, height: int = HEIGHT, channels: int = 2) -> xr.DataArray: + y, x = np.mgrid[0:height, 0:width] + base = ((np.sin(x / 17.0) * np.cos(y / 11.0) + 1) * 1000).astype(np.uint16) + stack = np.stack([base * (i + 1) for i in range(channels)]) + return xr.DataArray(stack, dims=("c", "y", "x"), coords={"c": [f"ch{i}" for i in range(channels)]}) + + +@pytest.fixture +def written(tmp_path: Path) -> tuple[Path, dict]: + out = tmp_path / "dapi" + manifest = write_webp_pyramid(_image(), out, tile_size=TILE, source_element="morphology") + return out, manifest + + +def _tile(directory: Path, manifest: dict, zoom: int, tx: int, ty: int): + zi = manifest["zoom_info"][str(zoom)] + rg = zi["row_group_offset"] + tx * zi["num_tiles_y"] + ty + file_index, local = divmod(rg, manifest["max_row_groups_per_file"]) + return pq.ParquetFile(directory / manifest["files"][file_index]).read_row_group(local) + + +# -- zoom numbering ----------------------------------------------------------- + + +def test_max_zoom_follows_deepzoom(written: tuple[Path, dict]) -> None: + _, manifest = written + assert manifest["max_zoom"] == math.ceil(math.log2(max(WIDTH, HEIGHT))) + + +def test_every_level_down_to_zero_exists(written: tuple[Path, dict]) -> None: + """vips dzsave emits levels 0..max; matching it keeps the two interchangeable.""" + _, manifest = written + assert sorted(int(z) for z in manifest["zoom_info"]) == list(range(manifest["max_zoom"] + 1)) + + +def test_level_dimensions_halve(written: tuple[Path, dict]) -> None: + _, manifest = written + max_zoom = manifest["max_zoom"] + for zoom in range(max_zoom, 0, -1): + w = math.ceil(WIDTH / 2 ** (max_zoom - zoom)) + h = math.ceil(HEIGHT / 2 ** (max_zoom - zoom)) + zi = manifest["zoom_info"][str(zoom)] + assert zi["num_tiles_x"] == max(1, math.ceil(w / TILE)), f"zoom {zoom}" + assert zi["num_tiles_y"] == max(1, math.ceil(h / TILE)), f"zoom {zoom}" + + +def test_row_group_offsets_are_cumulative(written: tuple[Path, dict]) -> None: + _, manifest = written + running = 0 + for zoom in sorted(manifest["zoom_info"], key=int): + zi = manifest["zoom_info"][zoom] + assert zi["row_group_offset"] == running + assert zi["num_tiles"] == zi["num_tiles_x"] * zi["num_tiles_y"] + running += zi["num_tiles"] + assert running == manifest["total_row_groups"] + + +def test_index_formula_addresses_the_right_tile(written: tuple[Path, dict]) -> None: + """Mirror of ImageRowGroupReader.computeRowGroupIndex.""" + directory, manifest = written + for zoom_s, zi in manifest["zoom_info"].items(): + for tx in range(zi["num_tiles_x"]): + for ty in range(zi["num_tiles_y"]): + t = _tile(directory, manifest, int(zoom_s), tx, ty) + assert (t["zoom"][0].as_py(), t["tile_x"][0].as_py(), t["tile_y"][0].as_py()) == ( + int(zoom_s), + tx, + ty, + ) + + +# -- payload ------------------------------------------------------------------ + + +def test_tiles_decode_as_webp(written: tuple[Path, dict]) -> None: + from PIL import Image + + directory, manifest = written + t = _tile(directory, manifest, manifest["max_zoom"], 0, 0) + img = Image.open(io.BytesIO(t["image_data"][0].as_py())) + assert img.format == "WEBP" + assert img.size == (TILE, TILE) + + +def test_edge_tile_is_cropped_not_padded(written: tuple[Path, dict]) -> None: + from PIL import Image + + directory, manifest = written + max_zoom = manifest["max_zoom"] + zi = manifest["zoom_info"][str(max_zoom)] + t = _tile(directory, manifest, max_zoom, zi["num_tiles_x"] - 1, 0) + img = Image.open(io.BytesIO(t["image_data"][0].as_py())) + assert img.size[0] == WIDTH - (zi["num_tiles_x"] - 1) * TILE + + +def test_schema_matches_celldega_reader(written: tuple[Path, dict]) -> None: + directory, manifest = written + schema = pq.ParquetFile(directory / manifest["files"][0]).schema_arrow + assert schema.names == ["zoom", "tile_x", "tile_y", "image_data"] + meta = schema.metadata + assert json.loads(meta[b"zoom_info"]) == manifest["zoom_info"] + assert meta[b"storage_mode"] == b"row_groups_image_chunked" + + +def test_one_tile_per_row_group(written: tuple[Path, dict]) -> None: + directory, manifest = written + for f in manifest["files"]: + md = pq.ParquetFile(directory / f).metadata + assert all(md.row_group(i).num_rows == 1 for i in range(md.num_row_groups)) + + +# -- channels and windowing --------------------------------------------------- + + +def test_channel_can_be_selected_by_name(tmp_path: Path) -> None: + a = write_webp_pyramid(_image(), tmp_path / "a", tile_size=TILE, channel="ch0") + b = write_webp_pyramid(_image(), tmp_path / "b", tile_size=TILE, channel="ch1") + # ch1 is twice ch0, so its auto window differs. + assert b["display_max"] > a["display_max"] + + +def test_explicit_window_is_recorded_and_used(tmp_path: Path) -> None: + m = write_webp_pyramid(_image(), tmp_path / "w", tile_size=TILE, display_min=100, display_max=900) + assert (m["display_min"], m["display_max"]) == (100.0, 900.0) + + +def test_degenerate_window_does_not_divide_by_zero(tmp_path: Path) -> None: + flat = xr.DataArray(np.full((1, 200, 200), 7, dtype=np.uint16), dims=("c", "y", "x")) + m = write_webp_pyramid(flat, tmp_path / "flat", tile_size=TILE) + assert m["display_max"] > m["display_min"] + + +def test_multiscale_input_uses_full_resolution(tmp_path: Path) -> None: + """A DataTree must be tiled from scale0, not from a downsampled level.""" + from xarray import DataTree + + full = _image(400, 200, channels=1) + half = full.isel(y=slice(None, None, 2), x=slice(None, None, 2)) + tree = DataTree.from_dict({"scale0": full.to_dataset(name="image"), "scale1": half.to_dataset(name="image")}) + m = write_webp_pyramid(tree, tmp_path / "ms", tile_size=TILE) + assert (m["source_width"], m["source_height"]) == (400, 200) + + +# -- metadata and safety ------------------------------------------------------ + + +def test_source_metadata_is_recorded_for_invalidation(written: tuple[Path, dict]) -> None: + _, m = written + assert m["source_element"] == "morphology" + assert (m["source_width"], m["source_height"]) == (WIDTH, HEIGHT) + assert m["source_dtype"] == "uint16" + assert m["downsampling"] == "box-2x2-mean" + assert m["image_format"] == ".webp" + + +def test_lossless_is_recorded(tmp_path: Path) -> None: + m = write_webp_pyramid(_image(), tmp_path / "ll", tile_size=TILE, lossless=True) + assert m["webp_lossless"] is True + assert m["webp_quality"] is None + + +def test_overwrite_guard(written: tuple[Path, dict]) -> None: + directory, _ = written + with pytest.raises(FileExistsError): + write_webp_pyramid(_image(), directory, tile_size=TILE) + + +def test_failure_leaves_no_partial_output(tmp_path: Path) -> None: + out = tmp_path / "boom" + bad = xr.DataArray(np.zeros((2, 2, 10, 10), dtype=np.uint16), dims=("t", "c", "y", "x")) + with pytest.raises(ValueError, match="expected a 2D plane"): + write_webp_pyramid(bad, out, tile_size=TILE) + assert not out.exists() + assert not out.with_name(out.name + ".tmp").exists() From bca1d84f35191c176835aa4abe72b396c1564aa4 Mon Sep 17 00:00:00 2001 From: Nicolas Fernandez Date: Sun, 6 Sep 2026 16:25:02 -0400 Subject: [PATCH 09/17] feat(experimental): stream the points rewrite so tiling survives limited memory Grouping rows by tile is a global sort -- a row at the end of the input can belong to the first tile -- so streaming the read alone is not enough. The streaming path makes two passes: stream the element one partition at a time and spill each row into a temporary file chosen by its destination chunk file, then sort each spill file independently and write its chunk. Peak memory scales with one partition plus one spill file rather than the dataset. Enabled automatically for partitioned (dask) elements; single-partition elements keep the single-pass path. A test asserts the two produce identical row groups, so they cannot silently diverge. Chunks pin the categorical dictionary to the element's full category list, since partitions observing different feature subsets would otherwise convert to incompatible Arrow dictionary types and could not be written to one file. Measured on Xenium Prime human skin (74,011,892 transcripts, 5,006 genes, 9.2x the pancreas): 94s to tile at 3.0 GB peak RSS, on a 17 GB machine where the in-memory path was estimated at 17-20 GB. 14,022 row groups across 36 files, all rows preserved, manifest validates, read_zarr returns the full element. Co-Authored-By: Claude Opus 5 (1M context) --- .../experimental/points_parquet.py | 336 ++++++++++++++---- tests/test_points_parquet.py | 57 +++ 2 files changed, 329 insertions(+), 64 deletions(-) diff --git a/src/spatialdata_io/experimental/points_parquet.py b/src/spatialdata_io/experimental/points_parquet.py index b68ce458..c57085e8 100644 --- a/src/spatialdata_io/experimental/points_parquet.py +++ b/src/spatialdata_io/experimental/points_parquet.py @@ -129,6 +129,93 @@ def _interleaved_positions(px: NDArray[np.uint32], py: NDArray[np.uint32]) -> pa return pa.FixedSizeListArray.from_arrays(pa.array(flat), 2) +#: Internal column carrying the tile assignment through the streaming spill files. +_TILE_ID = "__tile_id" + + +def _prepare_table( + df: pd.DataFrame, + *, + transform: DisplayTransform, + catalog: FeatureCatalog, + feature_key: str, + grid: RegularGrid, + categories: Any | None, +) -> tuple[pa.Table, NDArray[np.int64]]: + """Add the render columns to one chunk and return it with its tile assignment. + + ``categories`` pins the categorical dictionary so that every chunk converts to an + identical Arrow schema; without it, partitions observing different feature subsets + would produce incompatible dictionary types and could not be written to one file. + """ + px, py = _to_display_pixels(df["x"].to_numpy(), df["y"].to_numpy(), transform) + tile_ids = grid.assign(px, py) + codes = catalog.encode(df[feature_key]) + + # Re-tiling an already-tiled element must replace the render columns, not append + # duplicates: a duplicated name makes projected reads fail ("Multiple matches for + # FieldRef"), and a pandas round-trip degrades fixed_size_list to a variable list. + stale = [c for c in (POSITION_COLUMN, FEATURE_COLUMN, _TILE_ID) if c in df.columns] + if df.attrs or stale or categories is not None: + df = df.copy(deep=False) + # The transform is not JSON-serializable; spatialdata's own points writer drops it + # the same way. It is persisted in the element's zarr attributes, not the parquet. + df.attrs = {} + if stale: + df = df.drop(columns=stale) + if categories is not None and isinstance(df[feature_key].dtype, pd.CategoricalDtype): + df[feature_key] = df[feature_key].cat.set_categories(categories) + + table = pa.Table.from_pandas(df, preserve_index=True) + table = table.append_column(POSITION_COLUMN, _interleaved_positions(px, py)) + table = table.append_column(FEATURE_COLUMN, pa.array(codes)) + return table, tile_ids + + +def _sorted_by_tile(table: pa.Table, tile_ids: NDArray[np.int64]) -> tuple[pa.Table, NDArray[np.int64]]: + """Group rows by tile. The sort is stable, so the rewrite is deterministic.""" + order = np.argsort(tile_ids, kind="stable") + return table.take(pa.array(order)), tile_ids[order] + + +def _write_tile_row_groups( + writer: pq.ParquetWriter, + table: pa.Table, + sorted_tile_ids: NDArray[np.int64], + tile_range: range, + schema: pa.Schema, +) -> None: + """Write one row group per tile in ``tile_range``, empty tiles included. + + Empty tiles must still occupy a row group, since that is what makes + ``row_group_index == tile_id`` hold without a lookup table. + """ + offsets = np.searchsorted(sorted_tile_ids, np.array([*tile_range, tile_range.stop]), side="left") + for i in range(len(tile_range)): + start, end = int(offsets[i]), int(offsets[i + 1]) + writer.write_table(table.slice(start, end - start) if end > start else schema.empty_table()) + + +def _iter_chunks(points: Any) -> Any: + """Yield the element one partition at a time, or once if it is already in memory.""" + if hasattr(points, "npartitions"): + for i in range(points.npartitions): + yield points.partitions[i].compute() + else: + yield points + + +def _known_categories(points: Any, feature_key: str) -> Any | None: + """Return the element's full category list, so every chunk shares one dictionary.""" + col = points[feature_key] + if not hasattr(col, "cat"): + return None + try: + return list(col.cat.categories) + except (NotImplementedError, AttributeError): + return list(col.cat.as_known().cat.categories) + + def write_points_regular_grid( points: Any, output_dir: str | Path, @@ -141,6 +228,7 @@ def write_points_regular_grid( tile_size_px: float = 250.0, max_row_groups_per_file: int = DEFAULT_MAX_ROW_GROUPS_PER_FILE, compression: str = "zstd", + streaming: bool | None = None, overwrite: bool = False, ) -> dict[str, Any]: """Write a Points element as regular-grid row groups. @@ -178,53 +266,48 @@ def write_points_regular_grid( if output_dir.exists() and not overwrite: raise FileExistsError(f"{output_dir} exists; pass overwrite=True to replace it") - df = points.compute() if hasattr(points, "compute") else points - if not isinstance(df, pd.DataFrame): - raise TypeError(f"expected a DataFrame, got {type(df).__name__}") - if feature_key not in df.columns: - raise ValueError(f"feature column {feature_key!r} not found; have {list(df.columns)}") + columns = list(points.columns) + if feature_key not in columns: + raise ValueError(f"feature column {feature_key!r} not found; have {columns}") for axis in ("x", "y"): - if axis not in df.columns: - raise ValueError(f"points element has no {axis!r} column; have {list(df.columns)}") + if axis not in columns: + raise ValueError(f"points element has no {axis!r} column; have {columns}") + + partitioned = hasattr(points, "npartitions") and points.npartitions > 1 + if streaming is None: + streaming = partitioned + if streaming and not partitioned: + raise ValueError( + "streaming requires a partitioned (dask) points element; an in-memory frame cannot be read incrementally" + ) # When called as a SpatialData ``points_writer`` hook the element arrives with its # transformations already stripped from attrs, so the caller must supply the transform. transform = display_transform or DisplayTransform.from_element(points, coordinate_system) - px, py = _to_display_pixels(df["x"].to_numpy(), df["y"].to_numpy(), transform) if grid is None: - grid = RegularGrid.from_bounds(0, 0, float(px.max()), float(py.max()), tile_size_px) - - tile_ids = grid.assign(px, py) - codes = catalog.encode(df[feature_key]) - - # Keep every canonical column and index; append the two render columns. - # The transform lives in .attrs and is not JSON-serializable, so drop it before the - # Arrow conversion exactly as spatialdata's own points writer does -- it is persisted - # in the element's zarr attributes, not in the parquet file. - stale = [c for c in (POSITION_COLUMN, FEATURE_COLUMN) if c in df.columns] - if df.attrs or stale: - df = df.copy(deep=False) - df.attrs = {} - # Re-running the optimizer on an already-optimized element must replace the render - # columns, not append duplicates. A duplicated name makes the file unreadable by - # column projection ("Multiple matches for FieldRef"), and pandas round-trips - # fixed_size_list back as a variable-length list, so the stale copy is also the - # wrong Arrow type. Both are recomputed below from the canonical coordinates. - if stale: - df = df.drop(columns=stale) - table = pa.Table.from_pandas(df, preserve_index=True) - table = table.append_column(POSITION_COLUMN, _interleaved_positions(px, py)) - table = table.append_column(FEATURE_COLUMN, pa.array(codes)) + grid = _derive_grid(points, transform, tile_size_px) + + if streaming: + return _write_streaming( + points, + output_dir, + catalog=catalog, + grid=grid, + transform=transform, + feature_key=feature_key, + max_row_groups_per_file=max_row_groups_per_file, + compression=compression, + ) - # Stable sort keeps the original relative order inside a tile, so the rewrite is - # deterministic and diffable. - order = np.argsort(tile_ids, kind="stable") - table = table.take(pa.array(order)) - sorted_tile_ids = tile_ids[order] + df = points.compute() if hasattr(points, "compute") else points + if not isinstance(df, pd.DataFrame): + raise TypeError(f"expected a DataFrame, got {type(df).__name__}") - # Row-group boundaries: offsets[t]..offsets[t+1] is tile t's slice. - offsets = np.searchsorted(sorted_tile_ids, np.arange(grid.num_tiles + 1), side="left") + table, tile_ids = _prepare_table( + df, transform=transform, catalog=catalog, feature_key=feature_key, grid=grid, categories=None + ) + table, sorted_tile_ids = _sorted_by_tile(table, tile_ids) staging = output_dir.with_name(output_dir.name + ".tmp") if staging.exists(): @@ -243,31 +326,19 @@ def write_points_regular_grid( ) try: - writer: pq.ParquetWriter | None = None - current_file = -1 - for tile_id in range(grid.num_tiles): - file_index, _ = grid.chunk_location(tile_id, max_row_groups_per_file) - if file_index != current_file: - if writer is not None: - writer.close() - writer = pq.ParquetWriter( - staging / filenames[file_index], - schema, - compression=compression, - # Statistics are dead weight here: the tile formula is the spatial - # index, so no client ever consults per-column-chunk min/max, and - # they inflate the footer the browser must download up front. - write_statistics=False, - ) - current_file = file_index - - start, end = int(offsets[tile_id]), int(offsets[tile_id + 1]) - assert writer is not None - # Empty tiles are written as zero-row row groups so that - # row_group_index == tile_id holds without a lookup table. - writer.write_table(table.slice(start, end - start) if end > start else schema.empty_table()) - if writer is not None: - writer.close() + for file_index, name in enumerate(filenames): + lo = file_index * max_row_groups_per_file + tile_range = range(lo, min(lo + max_row_groups_per_file, grid.num_tiles)) + with pq.ParquetWriter( + staging / name, + schema, + compression=compression, + # Statistics are dead weight here: the tile formula is the spatial index, + # so no client consults per-column-chunk min/max, and they inflate the + # footer the browser must download before its first read. + write_statistics=False, + ) as writer: + _write_tile_row_groups(writer, table, sorted_tile_ids, tile_range, schema) except BaseException: shutil.rmtree(staging, ignore_errors=True) raise @@ -276,6 +347,30 @@ def write_points_regular_grid( shutil.rmtree(output_dir) staging.rename(output_dir) + return _manifest_fragment( + output_dir, filenames, grid, transform, catalog, max_row_groups_per_file, int(table.num_rows) + ) + + +def _derive_grid(points: Any, transform: DisplayTransform, tile_size_px: float) -> RegularGrid: + """Find the grid covering the element, reading only the coordinate columns.""" + xmax = points["x"].max() + ymax = points["y"].max() + if hasattr(xmax, "compute"): + xmax, ymax = xmax.compute(), ymax.compute() + px, py = transform.apply(np.array([float(xmax)]), np.array([float(ymax)])) + return RegularGrid.from_bounds(0, 0, float(np.rint(px[0])), float(np.rint(py[0])), tile_size_px) + + +def _manifest_fragment( + output_dir: Path, + filenames: list[str], + grid: RegularGrid, + transform: DisplayTransform, + catalog: FeatureCatalog, + max_row_groups_per_file: int, + n_rows: int, +) -> dict[str, Any]: return { "directory": str(output_dir.name), "files": filenames, @@ -289,8 +384,121 @@ def write_points_regular_grid( # Projected by the client, so canonical coordinates, ids and QC columns are never # decoded or transferred during ordinary rendering. "columns": [POSITION_COLUMN, FEATURE_COLUMN], - "n_rows": int(table.num_rows), + "n_rows": n_rows, "tile_grid": grid.to_manifest_dict(), "display_transform": transform.to_manifest_dict(), "feature_catalog": catalog.to_manifest_dict(), } + + +def _write_streaming( + points: Any, + output_dir: Path, + *, + catalog: FeatureCatalog, + grid: RegularGrid, + transform: DisplayTransform, + feature_key: str, + max_row_groups_per_file: int, + compression: str, +) -> dict[str, Any]: + """Write the tiled output without holding the whole element in memory. + + Grouping rows by tile is a global sort -- a row at the end of the input can belong to + the first tile -- so streaming the read alone is not enough. This makes two passes: + + 1. Stream the input a partition at a time, and spill each row into a temporary file + chosen by its *destination chunk file*. + 2. Sort each spill file independently and write its chunk. + + Peak memory is then one input partition plus one spill file, rather than the dataset. + """ + filenames = grid.chunk_filenames(max_row_groups_per_file) + categories = _known_categories(points, feature_key) + + staging = output_dir.with_name(output_dir.name + ".tmp") + spill = output_dir.with_name(output_dir.name + ".spill") + for path in (staging, spill): + if path.exists(): + shutil.rmtree(path) + staging.mkdir(parents=True) + spill.mkdir(parents=True) + + schema: pa.Schema | None = None + n_rows = 0 + try: + # -- pass 1: spill by destination file --------------------------------- + spill_writers: dict[int, pq.ParquetWriter] = {} + for chunk in _iter_chunks(points): + if len(chunk) == 0: + continue + table, tile_ids = _prepare_table( + chunk, + transform=transform, + catalog=catalog, + feature_key=feature_key, + grid=grid, + categories=categories, + ) + n_rows += table.num_rows + if schema is None: + schema = table.schema + table = table.append_column(_TILE_ID, pa.array(tile_ids)) + + buckets = tile_ids // max_row_groups_per_file + for bucket in np.unique(buckets): + rows = np.flatnonzero(buckets == bucket) + part = table.take(pa.array(rows)) + if bucket not in spill_writers: + spill_writers[int(bucket)] = pq.ParquetWriter( + spill / f"{int(bucket)}.parquet", table.schema, compression="zstd", write_statistics=False + ) + spill_writers[int(bucket)].write_table(part) + for writer in spill_writers.values(): + writer.close() + + if schema is None: + raise ValueError("points element is empty; nothing to tile") + schema = schema.with_metadata( + { + **(schema.metadata or {}), + b"profile": b"celldega_regular_grid_v1", + b"storage_mode": b"row_groups_chunked", + b"max_row_groups_per_file": str(max_row_groups_per_file).encode(), + b"tile_grid": json.dumps(grid.to_manifest_dict()).encode(), + } + ) + + # -- pass 2: sort each spill file and write its chunk ------------------- + for file_index, name in enumerate(filenames): + lo = file_index * max_row_groups_per_file + tile_range = range(lo, min(lo + max_row_groups_per_file, grid.num_tiles)) + spill_path = spill / f"{file_index}.parquet" + + if spill_path.exists(): + table = pq.read_table(spill_path) + tile_ids = table[_TILE_ID].to_numpy() + table = table.drop_columns([_TILE_ID]).cast(schema) + table, sorted_tile_ids = _sorted_by_tile(table, tile_ids) + else: + # No row landed in this chunk's tiles; it is still written, as all-empty + # row groups, so the global row-group numbering stays contiguous. + table, sorted_tile_ids = schema.empty_table(), np.empty(0, dtype=np.int64) + + with pq.ParquetWriter(staging / name, schema, compression=compression, write_statistics=False) as writer: + _write_tile_row_groups(writer, table, sorted_tile_ids, tile_range, schema) + + del table + spill_path.unlink(missing_ok=True) + except BaseException: + shutil.rmtree(staging, ignore_errors=True) + shutil.rmtree(spill, ignore_errors=True) + raise + finally: + shutil.rmtree(spill, ignore_errors=True) + + if output_dir.exists(): + shutil.rmtree(output_dir) + staging.rename(output_dir) + + return _manifest_fragment(output_dir, filenames, grid, transform, catalog, max_row_groups_per_file, n_rows) diff --git a/tests/test_points_parquet.py b/tests/test_points_parquet.py index d8105cbf..278944a6 100644 --- a/tests/test_points_parquet.py +++ b/tests/test_points_parquet.py @@ -307,3 +307,60 @@ def test_rewrite_is_idempotent(tmp_path: Path, points: pd.DataFrame, catalog: Fe assert second.schema.field(POSITION_COLUMN).type == pa.list_(pa.uint32(), 2) assert second.num_rows == first.num_rows assert second[POSITION_COLUMN].to_pylist() == first[POSITION_COLUMN].to_pylist() + + +def test_streaming_matches_in_memory(tmp_path: Path, points: pd.DataFrame, catalog: FeatureCatalog) -> None: + """The two write paths must be interchangeable, or large datasets would diverge. + + Grouping by tile is a global sort, so the streaming path spills rows into per-output + -file buckets and sorts each independently. That must land every row in the same row + group as the single-pass path. + """ + import dask.dataframe as dd + from spatialdata.transformations import Scale, set_transformation + + set_transformation(points, Scale([2.0, 2.0], axes=("x", "y")), "global") + in_memory = tmp_path / "mem.parquet" + m1 = write_points_regular_grid(points, in_memory, catalog=catalog, grid=GRID, streaming=False) + + # Several partitions, so the spill path is genuinely exercised. + chunked = dd.from_pandas(points.compute(), npartitions=3) + chunked.attrs["transform"] = {"global": Scale([2.0, 2.0], axes=("x", "y"))} + streamed = tmp_path / "stream.parquet" + m2 = write_points_regular_grid( + chunked, streamed, catalog=catalog, grid=GRID, streaming=True, max_row_groups_per_file=2 + ) + + assert m1["total_row_groups"] == m2["total_row_groups"] == GRID.num_tiles + assert m1["n_rows"] == m2["n_rows"] == len(POINTS_SPEC) + + # Every tile must hold exactly the same points in both. + for tile_id in range(GRID.num_tiles): + fi1, lo1 = GRID.chunk_location(tile_id, m1["max_row_groups_per_file"]) + fi2, lo2 = GRID.chunk_location(tile_id, m2["max_row_groups_per_file"]) + a = pq.ParquetFile(in_memory / m1["files"][fi1]).read_row_group(lo1, columns=[POSITION_COLUMN]) + b = pq.ParquetFile(streamed / m2["files"][fi2]).read_row_group(lo2, columns=[POSITION_COLUMN]) + assert a[POSITION_COLUMN].to_pylist() == b[POSITION_COLUMN].to_pylist(), f"tile {tile_id}" + + +def test_streaming_preserves_canonical_columns(tmp_path: Path, points: pd.DataFrame, catalog: FeatureCatalog) -> None: + import dask.dataframe as dd + from spatialdata.transformations import Scale + + original = points.compute() + chunked = dd.from_pandas(original, npartitions=3) + chunked.attrs["transform"] = {"global": Scale([2.0, 2.0], axes=("x", "y"))} + out = tmp_path / "s.parquet" + manifest = write_points_regular_grid(chunked, out, catalog=catalog, grid=GRID, streaming=True) + + got = _read_all(out, manifest).to_pandas().set_index("transcript_id").loc[original["transcript_id"].to_numpy()] + for col in ("x", "y", "z", "cell_id", "qv"): + np.testing.assert_array_equal(got[col].to_numpy(), original[col].to_numpy(), err_msg=col) + assert list(got["feature_name"].astype(str)) == list(original["feature_name"].astype(str)) + + +def test_streaming_requires_a_partitioned_element( + tmp_path: Path, points: pd.DataFrame, catalog: FeatureCatalog +) -> None: + with pytest.raises(ValueError, match="streaming requires a partitioned"): + write_points_regular_grid(points.compute(), tmp_path / "x.parquet", catalog=catalog, grid=GRID, streaming=True) From acf54b6aef0efbe59d3827209e14a1e7238e4201 Mon Sep 17 00:00:00 2001 From: Nicolas Fernandez Date: Sun, 6 Sep 2026 17:53:27 -0400 Subject: [PATCH 10/17] refactor(experimental): rename profile to grid_files_v1 and add the fixed-path assets Renames the profile from celldega_regular_grid_v1 to grid_files_v1. Naming it after one client contradicted the viewer-independent intent; nothing about the layout is Celldega-specific. Adds the files a client reads at fixed paths rather than through the manifest, which is what lets the profile directory stand in for a DegaFiles root so no client needs to know it is looking at a SpatialData store: - cell_metadata.parquet: per-cell centroids and names. Cells are the overview representation, so every centroid is needed up front. Centroids are not stored anywhere in a SpatialData store (the Xenium reader puts no x/y_centroid in obs), so they are derived from the shape geometry -- the same centroids already computed for tile assignment. Row order matches the table's obs order, so a cell's position here is its cell_code. - micron_to_image_transform.csv: the real micron-to-pixel affine. Coordinates are already in display pixels so this places nothing, but it drives the scale bar, and writing identity would make the scale bar wrong by the pixel size. - cell_clusters/: placeholder single-group clustering, since SpatialData does not require one and the Xenium reader does not load one. Without the file the viewer's category machinery fails on a 404. Also drops the design doc from the branch (it moves to the PR discussion) and removes supports_points_writer_hook, which was unused and untested. Co-Authored-By: Claude Opus 5 (1M context) --- docs/design/regular_grid_tiled_access.md | 426 ------------------ .../experimental/cbg_parquet.py | 2 +- src/spatialdata_io/experimental/manifest.py | 9 +- .../experimental/points_parquet.py | 4 +- .../experimental/regular_grid.py | 2 +- .../experimental/shapes_parquet.py | 68 ++- .../experimental/tiled_access.py | 97 +++- .../experimental/webp_parquet.py | 2 +- 8 files changed, 161 insertions(+), 449 deletions(-) delete mode 100644 docs/design/regular_grid_tiled_access.md diff --git a/docs/design/regular_grid_tiled_access.md b/docs/design/regular_grid_tiled_access.md deleted file mode 100644 index cb957eac..00000000 --- a/docs/design/regular_grid_tiled_access.md +++ /dev/null @@ -1,426 +0,0 @@ -# Regular-grid tiled access profile (`celldega_regular_grid_v1`) - -**Status:** experimental · **Profile version:** 0.1.0 - -A viewer-independent protocol for fetching spatially-local subsets of a SpatialData store -over HTTP range requests, without downloading whole files and without consulting Parquet -statistics. - -Celldega is the first reference client. Nothing in this document is Celldega-specific; -another viewer (Vitessce, SpatialData.js, napari) could implement it from this text alone. - ---- - -## 1. Motivation and scope - -A viewer showing a 34,000 × 14,000 pixel tissue with 8 million transcripts must fetch only -what is on screen. Two things have to be true: - -1. The client can compute **which bytes it needs** from the viewport alone — no index - download, no metadata probing, no statistics. -2. Those bytes are **already in the form the GPU wants** — no per-point object - construction, no coordinate zipping, no WKB parsing. - -This profile achieves both by reordering rows into a deterministic grid of Parquet row -groups and adding a small number of render-oriented columns. - -**In scope:** transcript points, cell polygons, gene-major expression, image tiles, -and the manifest that describes them. - -**Out of scope:** clustering, linked views, annotation, neighbourhood analysis, and any -other viewer feature. This is a data-access protocol. - -### Design invariants - -- **The profile is opt-in.** Default SpatialData write behaviour is unchanged. -- **Canonical data is authoritative and preserved.** Render columns are additions; they - never replace canonical coordinates, identifiers, geometries or annotations. -- **A store carrying the profile is still an ordinary SpatialData store.** - `spatialdata.read_zarr()` works unchanged; a client that does not understand the profile - ignores the extra columns and the manifest. - ---- - -## 2. Coordinate system - -All profile coordinates are **level-0 pixels of a chosen reference image**, referred to -here as *display pixel space*. - -The mapping from canonical element coordinates to display pixels is the element's -SpatialData affine transformation into a named coordinate system (`global` by default). -It is recorded in the manifest: - -```json -"display_transform": { - "coordinate_space": "image-pixel", - "coordinate_system": "global", - "affine_matrix": [[4.70588235, 0.0, 0.0], [0.0, 4.70588235, 0.0]], - "rounding": "nearest" -} -``` - -Producers **MUST** apply the affine in float64 and round half-to-even (`rint`). Applying -it in float32 loses pixel accuracy on large images, and in some environments a float32 -array multiplied by a scalar stays float32. - -Display coordinates **MUST** be non-negative and **MUST** fit the declared integer dtype. -A producer encountering values outside that range **MUST** fail rather than clamp: an -out-of-range coordinate indicates a mismatched transform, not a rounding artefact. - ---- - -## 3. The grid - -A non-overlapping regular square grid over display pixel space, defined by five numbers: - -| field | meaning | -|---|---| -| `x_min`, `y_min` | grid origin, display pixels | -| `tile_size` | tile edge length, display pixels | -| `num_tiles_x`, `num_tiles_y` | grid dimensions | - -### Tile assignment - -``` -tile_x = floor((x_px - x_min) / tile_size) -tile_y = floor((y_px - y_min) / tile_size) -``` - -Tile bounds are **half-open** `[min, max)`: a coordinate lying exactly on an internal -boundary belongs to the **upper** tile. - -The single exception is the grid's outer edge. A coordinate equal to `x_max` or `y_max` -is clamped into the last tile, so that a point on the boundary of the dataset is not -lost. Coordinates beyond one tile past the extent are an error. - -### Tile numbering - -Tiles are numbered **x-major**: - -``` -tile_id = tile_x * num_tiles_y + tile_y -``` - -`tile_id` ranges over `[0, num_tiles_x * num_tiles_y)`. - -### Choosing `tile_size` - -`tile_size` is a tuning parameter, not a constant. It trades viewport granularity against -storage: smaller tiles fetch less off-screen data but fragment the file into more, -individually-compressed row groups. - -The recommended target is **roughly 20 cells per tile**, which is the granularity at -which a viewer fetches. On Xenium-density tissue this is about **250 display pixels**. -Measured on a Xenium human pancreas section (140,702 cells): - -| tile px | cells/tile | row groups | size vs untiled | -|---|---|---|---| -| 200 | 13.5 | 11,799 | +71% | -| **250** | **21.0** | **7,535** | **+60%** | -| 500 | 81.0 | 1,932 | +50% | - ---- - -## 4. Row groups and files - -### One tile, one row group - -Each logical tile is written as **exactly one Parquet row group**, at index `tile_id`. -Tiles containing no rows are written as **zero-row row groups**, not skipped. This is what -lets a client address a tile by formula with no lookup table. - -A conforming file therefore contains exactly `num_tiles_x * num_tiles_y` row groups across -all its parts. - -### Multi-part files - -Row groups are split across files: - -``` -file_index = tile_id // max_row_groups_per_file -local_row_group = tile_id % max_row_groups_per_file -``` - -`max_row_groups_per_file` defaults to **400**. - -Splitting is **required**, not cosmetic. A Parquet reader must fetch a file's entire -footer before reading any row group, and footer size grows with row-group count. On the -pancreas dataset a single 7,535-row-group file has a **7.4 MB footer**; split into 19 -files each footer is ~410 KB, and a client only fetches footers for files its viewport -actually touches. - -### File naming - -Chunk files are named `chunk_.parquet` with `` **zero-padded** to the width of the -largest index (`chunk_00.parquet` … `chunk_18.parquet`). - -Padding is required because consumers disagree about ordering: a client indexes the -manifest's `files` array by position, but tools that glob a directory sort -lexicographically, where `chunk_10` precedes `chunk_2`. Padding makes the two agree. - -> Existing Celldega DegaFiles use unpadded names. That is safe there because only the -> manifest-array consumer exists. Stores written under this profile use padded names. - -### Statistics - -Producers **SHOULD** write Parquet files with column statistics disabled. The tile formula -is the spatial index, so no conforming client reads column-chunk min/max, and statistics -inflate the footer the client must download before its first read. - -### Compression - -**zstd** is the recommended codec. On the pancreas dataset, snappy costs +38.5% over an -untiled store while zstd costs +4.9% for the same content, at no meaningful write cost. -Producers **MUST NOT** use a codec the target client cannot decode. - ---- - -## 5. Transcript points - -The canonical Points element gains two columns; every canonical column and the DataFrame -index are preserved. Physical row order changes (rows are grouped by tile), which is -permitted; rows **MUST NOT** be added, dropped or altered. - -### `display_xy` - -``` -fixed_size_list[2] -``` - -Integer display-pixel coordinates. The Arrow child buffer is therefore already -`[x0, y0, x1, y1, ...]`, directly usable as a deck.gl binary `getPosition` attribute. - -A client **MUST NOT** need to interleave separate x and y arrays. - -> A future revision may allow fixed-point sub-pixel coordinates -> (`stored = pixel * 16`, `scale = 0.0625`). Producers of v0.1.0 write integer pixels -> and declare `"scale": 1.0`. - -### `feature_code` - -``` -uint16 (or uint32 when the catalog exceeds 65535 entries) -``` - -An index into the feature catalog (§7). - ---- - -## 6. Cell polygons - -The canonical Shapes element gains two columns. The canonical geometry column and its -GeoParquet `geo` metadata are preserved, so the file remains readable by -`geopandas.read_parquet` and by SpatialData. - -### `display_geometry` - -``` -list[2]>> -``` - -Polygon → rings → interleaved integer pixel vertices. A client lifts `getPolygon` from the -flat coordinate child buffer and `startIndices` from the list offsets: - -``` -start_index[i] = ring_offsets[polygon_offsets[i]] -``` - -`display_geometry` is explicitly a **lossy display representation**: it holds the exterior -ring only, and for a MultiPolygon only the largest part. The canonical geometry is retained -alongside it and is authoritative. - -> **Interoperability warning.** GeoArrow permits `struct` coordinates as well as -> interleaved `fixed_size_list`. A client walking the nesting blindly will, on struct -> coordinates, obtain the `x` child alone and render wrong polygons with no error. Clients -> **MUST** verify the vertex level is a `FixedSizeList` before treating the buffer as -> interleaved. (`geopandas.to_parquet` emits struct coordinates, so this is reachable in -> practice.) - -### `cell_code` - -``` -uint32 -``` - -Positional index into the annotating table's `obs` order, so cells can be coloured from an -expression vector without a string join in the client. - -### Tile assignment - -A cell is assigned to **exactly one tile**, by its **centroid** in display pixel space. - -A polygon whose outline crosses into neighbouring tiles is **NOT** duplicated into them. -Duplication would inflate the file and make cell counts wrong. A client rendering a -viewport should expect polygons to overhang tile boundaries and, if it needs full coverage -at the edges, fetch one extra ring of tiles. - ---- - -## 7. Feature catalog - -An ordered vocabulary mapping feature names to `feature_code`, stored as -`meta_gene.parquet` with columns `name`, `feature_code`, `is_gene`. - -Ordering is normative: - -1. Codes `[0, n_genes)` are **genes**, in the annotating table's `var_names` order. -2. Codes `>= n_genes` are **non-gene features** (negative controls, unassigned codewords), - sorted for reproducibility. - -Because genes come first and in table order, **a gene's `feature_code` is also its CBG row -group** (§8). One integer addresses both a transcript's identity and its expression vector. - -Non-gene features **MUST** be retained and **MUST NOT** be folded into a gene; doing so -would fabricate expression. `n_genes` is recorded in the manifest so a client can tell the -two apart. - ---- - -## 8. Cell-by-gene expression - -Gene-major, one row group per gene, so selecting a gene fetches one row group and touches -no transcript data. - -Columns: - -| column | type | meaning | -|---|---|---| -| `cell_id` | uint32 | cell code (§6), not a string barcode | -| `expression` | float32 | non-zero value | -| `gene` | string | gene name | - -Zero values are omitted, including sparse *stored* zeros. - -Row group `i` holds catalog gene `i`, including genes with no expression (written as a -zero-row row group), preserving the `feature_code == row group` invariant. The explicit -mapping is written both in the manifest and in the Parquet schema metadata: - -``` -gene_to_row_group JSON object -num_genes integer -storage_mode "row_groups_cbg_chunked" -``` - -A client **MAY** use the mapping rather than assuming the identity. - -> This is a **transpose**, not a redundant copy. SpatialData tables are stored cell-major -> (CSR); assembling one gene's vector from them requires reading the whole matrix or doing -> one random access per cell. - ---- - -## 9. Image tiles - -*(Not implemented in v0.1.0; specified here for forward compatibility.)* - -Canonical OME-Zarr images are retained and authoritative. An optional derived WebP pyramid -may be stored as Parquet row groups with columns `zoom`, `tile_x`, `tile_y`, `image_data` -(encoded WebP bytes). - -Image tiles at zoom 0 **MUST** use the same level-0 pixel coordinate system as -`display_xy` and `display_geometry`. - -Per channel the manifest records: source image element, source dimensions and dtype, -reference pyramid level, tile size, per-zoom grid dimensions, display intensity min/max, -gamma, colour, downsampling method, and WebP lossless/lossy setting. - ---- - -## 10. The manifest - -A JSON document named `landscape_parameters.json`, conventionally at -`.zarr/visualization/celldega_regular_grid_v1/`. - -Paths inside it are **relative to the manifest's own directory**, so a client pointed at -that directory resolves into the store without knowing the zarr layout: - -```json -{ - "technology": "Xenium", - "use_row_groups": true, - "profile": "celldega_regular_grid_v1", - "profile_version": "0.1.0", - "tile_grid": { - "num_tiles_x": 137, "num_tiles_y": 55, "tile_size": 250.0, - "x_min": 0.0, "y_min": 0.0, "x_max": 34250.0, "y_max": 13750.0 - }, - "row_group_files": { - "transcripts": { - "directory": "../../points/transcripts/points.parquet", - "files": ["chunk_00.parquet", "..."], - "max_row_groups_per_file": 400, - "total_row_groups": 7535, - "position_column": "display_xy", - "feature_column": "feature_code", - "columns": ["display_xy", "feature_code"] - }, - "cell_segmentation": { - "directory": "../../shapes/cell_boundaries/shapes.parquet", - "geometry_column": "display_geometry", - "cell_id_column": "cell_code", - "columns": ["display_geometry", "cell_code"] - }, - "cbg": { "directory": "cbg", "gene_to_row_group": {} }, - "images": {} - }, - "image_info": [] -} -``` - -`columns` is the projection a client should request. Honouring it is what keeps canonical -coordinates, identifiers and QC columns off the wire during rendering. - -A client that does not recognise the profile-specific keys and finds no column names -declared **SHOULD** fall back to reading all columns. - ---- - -## 11. Transport requirements - -A conforming host **MUST** support: - -- **HTTP range requests**, answering `Range: bytes=a-b` with `206 Partial Content` and a - correct `Content-Range`. -- **Suffix ranges** (`Range: bytes=-8`), which is how a reader locates the Parquet footer. -- **CORS**, with `Access-Control-Allow-Origin` and `Access-Control-Expose-Headers` - including `Content-Range`, for browser clients on another origin. - -Hugging Face dataset `resolve/` URLs satisfy all three. - ---- - -## 12. Invalidation - -The profile is derived data and **MUST** be regenerated when any of the following change: - -- transcript rows are added, removed, or spatially moved; -- the feature catalog or `var_names` order changes; -- canonical cell identifiers or table row order change (invalidates `cell_code`); -- cell geometries change; -- the reference image, or the transform into display pixel space, changes; -- `tile_size`, the grid origin, or `max_row_groups_per_file` change; -- image intensity windowing or WebP settings change (images only). - -Changing a SpatialData transformation that does **not** affect the declared display -coordinate system does not require regeneration, but this **MUST** be verified rather than -assumed — compare the resulting affine against `display_transform.affine_matrix`. - -Producers **SHOULD** record a source fingerprint in `source` so staleness is detectable. - ---- - -## 13. Conformance checklist - -A producer conforms if: - -- [ ] total row groups equals `num_tiles_x * num_tiles_y`, empty tiles included -- [ ] row group index equals `tile_x * num_tiles_y + tile_y` -- [ ] every canonical row appears exactly once; index preserved -- [ ] canonical columns and geometries are bit-identical to the source -- [ ] `display_xy` is `fixed_size_list[2]` with an interleaved child buffer -- [ ] `display_geometry` vertices are `fixed_size_list`, not `struct` -- [ ] each cell appears exactly once, in its centroid's tile -- [ ] `feature_code` matches the catalog; genes precede non-genes -- [ ] CBG row group index equals `feature_code` for every gene -- [ ] the store still opens with `spatialdata.read_zarr()` -- [ ] the manifest validates and every declared file exists diff --git a/src/spatialdata_io/experimental/cbg_parquet.py b/src/spatialdata_io/experimental/cbg_parquet.py index a717674a..92d8096b 100644 --- a/src/spatialdata_io/experimental/cbg_parquet.py +++ b/src/spatialdata_io/experimental/cbg_parquet.py @@ -116,7 +116,7 @@ def write_cbg_row_groups( b"storage_mode": b"row_groups_cbg_chunked", b"num_genes": str(n_genes).encode(), b"max_row_groups_per_file": str(max_row_groups_per_file).encode(), - b"profile": b"celldega_regular_grid_v1", + b"profile": b"grid_files_v1", } ) diff --git a/src/spatialdata_io/experimental/manifest.py b/src/spatialdata_io/experimental/manifest.py index 778ebcd3..993a2e21 100644 --- a/src/spatialdata_io/experimental/manifest.py +++ b/src/spatialdata_io/experimental/manifest.py @@ -29,7 +29,7 @@ "write_manifest", ] -PROFILE_NAME = "celldega_regular_grid_v1" +PROFILE_NAME = "grid_files_v1" PROFILE_VERSION = "0.1.0" MANIFEST_FILENAME = "landscape_parameters.json" @@ -45,6 +45,7 @@ def build_manifest( image_info: list[dict[str, Any]] | None = None, image_dimensions: dict[str, Any] | None = None, max_pyramid_zoom: int | None = None, + fixed_path_assets: dict[str, Any] | None = None, source: dict[str, Any] | None = None, ) -> dict[str, Any]: """Assemble the profile manifest from the fragments returned by each writer. @@ -94,6 +95,12 @@ def build_manifest( "profile": PROFILE_NAME, "profile_version": PROFILE_VERSION, } + if fixed_path_assets: + # Files a client reads by convention rather than through row_group_files: + # cell_metadata.parquet, meta_gene.parquet, micron_to_image_transform.csv, + # cell_clusters/. Recorded so the profile is self-describing even though the + # client does not consult these entries to find them. + manifest["fixed_path_assets"] = fixed_path_assets if image_dimensions is not None: manifest["image_dimensions"] = image_dimensions if max_pyramid_zoom is not None: diff --git a/src/spatialdata_io/experimental/points_parquet.py b/src/spatialdata_io/experimental/points_parquet.py index c57085e8..2d756ae4 100644 --- a/src/spatialdata_io/experimental/points_parquet.py +++ b/src/spatialdata_io/experimental/points_parquet.py @@ -318,7 +318,7 @@ def write_points_regular_grid( schema = table.schema.with_metadata( { **(table.schema.metadata or {}), - b"profile": b"celldega_regular_grid_v1", + b"profile": b"grid_files_v1", b"storage_mode": b"row_groups_chunked", b"max_row_groups_per_file": str(max_row_groups_per_file).encode(), b"tile_grid": json.dumps(grid.to_manifest_dict()).encode(), @@ -462,7 +462,7 @@ def _write_streaming( schema = schema.with_metadata( { **(schema.metadata or {}), - b"profile": b"celldega_regular_grid_v1", + b"profile": b"grid_files_v1", b"storage_mode": b"row_groups_chunked", b"max_row_groups_per_file": str(max_row_groups_per_file).encode(), b"tile_grid": json.dumps(grid.to_manifest_dict()).encode(), diff --git a/src/spatialdata_io/experimental/regular_grid.py b/src/spatialdata_io/experimental/regular_grid.py index 4d19334d..4f82245b 100644 --- a/src/spatialdata_io/experimental/regular_grid.py +++ b/src/spatialdata_io/experimental/regular_grid.py @@ -1,7 +1,7 @@ """Deterministic regular-grid spatial tiling. This module defines the tile geometry and the tile -> row-group -> file numbering -used by the ``celldega_regular_grid_v1`` visualization profile. +used by the ``grid_files_v1`` visualization profile. The grid is a non-overlapping regular square grid in *display pixel* space (level-0 pixels of a chosen reference image). Given an origin, a tile size and grid dimensions, diff --git a/src/spatialdata_io/experimental/shapes_parquet.py b/src/spatialdata_io/experimental/shapes_parquet.py index 8054c604..654ea4a1 100644 --- a/src/spatialdata_io/experimental/shapes_parquet.py +++ b/src/spatialdata_io/experimental/shapes_parquet.py @@ -40,7 +40,12 @@ RegularGrid, ) -__all__ = ["GEOMETRY_COLUMN", "CELL_CODE_COLUMN", "write_shapes_regular_grid"] +__all__ = [ + "GEOMETRY_COLUMN", + "CELL_CODE_COLUMN", + "write_shapes_regular_grid", + "write_cell_metadata", +] #: Column holding the nested integer-pixel display polygons. GEOMETRY_COLUMN = "display_geometry" @@ -196,7 +201,7 @@ def write_shapes_regular_grid( schema = table.schema.with_metadata( { **(table.schema.metadata or {}), - b"profile": b"celldega_regular_grid_v1", + b"profile": b"grid_files_v1", b"storage_mode": b"row_groups_chunked", b"max_row_groups_per_file": str(max_row_groups_per_file).encode(), b"tile_grid": json.dumps(grid.to_manifest_dict()).encode(), @@ -256,3 +261,62 @@ def write_shapes_regular_grid( fragment["directory"] = output_path.name fragment["files"] = filenames return fragment + + +def write_cell_metadata( + shapes: Any, + output_path: str | Path, + *, + display_transform: DisplayTransform | None = None, + coordinate_system: str = "global", + cell_index: Any | None = None, +) -> dict[str, Any]: + """Write the per-cell centroid table used for the overview scatter layer. + + Cells, not sampled transcripts, are the overview representation, so a client needs + every cell's centroid up front. This is a single small file (a few MB for ~10^5 cells) + rather than a tiled one, because the whole set is wanted at once. + + Row order is significant: a client takes a cell's integer id from its *position* here, + so the order must match ``cell_index`` (the annotating table's ``obs`` order) and hence + the ``cell_code`` written into the tiled shapes and the CBG. + + The schema matches Celldega's ``cell_metadata.parquet``: ``name`` plus ``geometry`` as + ``list`` holding ``[x, y]`` in display pixels. + """ + output_path = Path(output_path) + transform = display_transform or DisplayTransform.from_element(shapes, coordinate_system) + + simple = _exterior_only(shapes.geometry) + centroids = shapely.centroid(simple) + cx, cy = _to_display_pixels(shapely.get_x(centroids), shapely.get_y(centroids), transform) + + order = np.arange(len(shapes)) + names = [str(k) for k in shapes.index] + if cell_index is not None: + position = {k: i for i, k in enumerate(shapes.index)} + missing = [k for k in cell_index if k not in position] + if missing: + raise ValueError( + f"{len(missing)} cell(s) in cell_index have no shape (e.g. {missing[:3]}); " + f"their centroid would be undefined." + ) + order = np.fromiter((position[k] for k in cell_index), dtype=np.int64, count=len(cell_index)) + names = [str(k) for k in cell_index] + + flat = np.empty(len(order) * 2, dtype=np.float64) + flat[0::2] = cx[order] + flat[1::2] = cy[order] + geometry = pa.ListArray.from_arrays(pa.array(np.arange(len(order) + 1, dtype=np.int32) * 2), pa.array(flat)) + + table = pa.table({"name": pa.array(names, pa.string()), "geometry": geometry}) + output_path.parent.mkdir(parents=True, exist_ok=True) + pq.write_table(table, output_path, compression="zstd") + + return { + "path": output_path.name, + "n_cells": int(table.num_rows), + "position_encoding": "list", + "coordinate_space": "image-pixel", + "order": "annotating table obs order (position == cell_code)", + } diff --git a/src/spatialdata_io/experimental/tiled_access.py b/src/spatialdata_io/experimental/tiled_access.py index b696278e..4c658cff 100644 --- a/src/spatialdata_io/experimental/tiled_access.py +++ b/src/spatialdata_io/experimental/tiled_access.py @@ -7,10 +7,8 @@ :func:`xenium_spatially_tiled` Read raw Xenium and write a tiled store in one go. -The one-shot path is internally ``read -> write -> tile`` because the tiling rewrites -*written* Parquet. When the installed SpatialData exposes the ``points_writer`` hook the -points element is written in its tiled form directly, avoiding writing it twice; otherwise -it falls back to writing normally and rewriting, which produces an identical store. +The one-shot path is internally ``read -> write -> tile``, because the tiling rewrites +*written* Parquet. Everything here is additive and opt-in. A store that has been tiled is still an ordinary SpatialData store: :func:`spatialdata.read_zarr` works unchanged, the canonical columns and @@ -20,7 +18,6 @@ from __future__ import annotations -import inspect import shutil from pathlib import Path from typing import Any @@ -43,21 +40,17 @@ DEFAULT_MAX_ROW_GROUPS_PER_FILE, RegularGrid, ) -from spatialdata_io.experimental.shapes_parquet import write_shapes_regular_grid +from spatialdata_io.experimental.shapes_parquet import ( + write_cell_metadata, + write_shapes_regular_grid, +) -__all__ = ["add_spatial_tiling", "xenium_spatially_tiled", "supports_points_writer_hook"] +__all__ = ["add_spatial_tiling", "xenium_spatially_tiled"] #: Directory inside the store holding derived (non-canonical) profile assets. PROFILE_DIR = "visualization" -def supports_points_writer_hook() -> bool: - """Whether the installed SpatialData accepts a ``points_writer`` in ``write()``.""" - from spatialdata import SpatialData - - return "points_writer" in inspect.signature(SpatialData.write).parameters - - def _grid_for(points: Any, transform: DisplayTransform, tile_size_px: float) -> RegularGrid: """Derive the grid covering the points element in display pixel space.""" x = points["x"].max().compute() if hasattr(points["x"].max(), "compute") else points["x"].max() @@ -170,15 +163,29 @@ def add_spatial_tiling( transcripts["directory"] = f"../../points/{points_element}/points.parquet" cell_segmentation = None + cell_metadata = None + cell_names: list[str] | None = None if shapes_element: if shapes_element not in sdata.shapes: raise ValueError(f"shapes element {shapes_element!r} not found; have {list(sdata.shapes)}") shapes = sdata.shapes[shapes_element] + shapes_transform = DisplayTransform.from_element(shapes, coordinate_system) + # Cells are the overview representation, so the client needs every centroid up + # front. Centroids are not stored anywhere in the SpatialData store (the Xenium + # reader puts no x/y_centroid in obs), so they are derived from the geometry -- + # the same centroids already computed for tile assignment. + cell_metadata = write_cell_metadata( + shapes, + profile_dir / "cell_metadata.parquet", + display_transform=shapes_transform, + cell_index=list(table.obs_names) if table is not None else None, + ) + cell_names = [str(k) for k in (table.obs_names if table is not None else shapes.index)] cell_segmentation = write_shapes_regular_grid( shapes, store / "shapes" / shapes_element / "shapes.parquet", grid=grid, - display_transform=DisplayTransform.from_element(shapes, coordinate_system), + display_transform=shapes_transform, cell_index=list(table.obs_names) if table is not None else None, max_row_groups_per_file=max_row_groups_per_file, compression=compression, @@ -232,6 +239,14 @@ def add_spatial_tiling( catalog.to_frame().to_parquet(profile_dir / "meta_gene.parquet", index=False) + # Files a client reads at fixed paths rather than through the manifest. Writing them + # is what lets the profile directory stand in for a DegaFiles root, so no client needs + # to know it is looking at a SpatialData store. + _write_micron_to_image_transform(profile_dir / "micron_to_image_transform.csv", transform) + cluster_info = None + if cell_names is not None: + cluster_info = _write_cell_clusters(profile_dir / "cell_clusters", cell_names, None) + manifest = build_manifest( grid=grid, technology=technology, @@ -242,6 +257,12 @@ def add_spatial_tiling( image_info=image_info, image_dimensions=image_dimensions, max_pyramid_zoom=max_pyramid_zoom, + fixed_path_assets={ + "meta_gene": "meta_gene.parquet", + "micron_to_image_transform": "micron_to_image_transform.csv", + **({"cell_metadata": cell_metadata} if cell_metadata else {}), + **({"cell_clusters": cluster_info} if cluster_info else {}), + }, source={ "store": store.name, "points_element": points_element, @@ -316,3 +337,49 @@ def xenium_spatially_tiled( compression=compression, **(tiling or {}), ) + + +def _write_micron_to_image_transform(path: Path, transform: DisplayTransform) -> None: + """Write the micron-to-image affine Celldega reads at a fixed path. + + Coordinates in the profile are already in display pixels, so this is not needed to + place anything. It is needed for the scale bar and any physical-units readout, so it + must be the real micron-to-pixel affine and not identity -- writing identity would + make the scale bar wrong by the pixel size. + """ + (a, b, c), (d, e, f) = transform.matrix + rows = [f"{a} {b} {c}", f"{d} {e} {f}", "0.0 0.0 1.0"] + path.write_text("\n".join(rows) + "\n") + + +def _write_cell_clusters(directory: Path, cell_names: list[str], clusters: Any | None) -> dict[str, Any]: + """Write the cluster assignment and palette Celldega reads at a fixed path. + + SpatialData does not require a clustering, and the Xenium reader does not load one, so + when none is supplied every cell is placed in a single group. That keeps the viewer's + category machinery working instead of failing on a missing file; it is a placeholder, + not a scientific result. + """ + import colorsys + + import pandas as pd + + directory.mkdir(parents=True, exist_ok=True) + if clusters is None: + labels = pd.Series(["unclustered"] * len(cell_names), index=cell_names, dtype=object) + else: + labels = pd.Series([str(v) for v in clusters], index=cell_names, dtype=object) + + counts = labels.value_counts() + palette = {} + for i, name in enumerate(counts.index): + r, g, b = colorsys.hsv_to_rgb((i * 0.618033988749895) % 1.0, 0.6, 0.9) + palette[name] = f"#{int(r * 255):02x}{int(g * 255):02x}{int(b * 255):02x}" + + pd.DataFrame({"cluster": labels}).to_parquet(directory / "cluster.parquet") + pd.DataFrame( + {"color": [palette[n] for n in counts.index], "count": counts.to_numpy()}, + index=list(counts.index), + ).to_parquet(directory / "meta_cluster.parquet") + + return {"n_clusters": int(counts.size), "placeholder": clusters is None} diff --git a/src/spatialdata_io/experimental/webp_parquet.py b/src/spatialdata_io/experimental/webp_parquet.py index 5c39c5b5..c8ea920c 100644 --- a/src/spatialdata_io/experimental/webp_parquet.py +++ b/src/spatialdata_io/experimental/webp_parquet.py @@ -234,7 +234,7 @@ def write_webp_pyramid( b"storage_mode": b"row_groups_image_chunked", b"max_row_groups_per_file": str(max_row_groups_per_file).encode(), b"tile_size": str(tile_size).encode(), - b"profile": b"celldega_regular_grid_v1", + b"profile": b"grid_files_v1", } ) From d133ced58bd11e0fd2b65d2cf566e09d35595edc Mon Sep 17 00:00:00 2001 From: Nicolas Fernandez Date: Sun, 6 Sep 2026 18:42:38 -0400 Subject: [PATCH 11/17] refactor(experimental): move render columns into their own parquet files The render columns no longer go into the canonical elements. Canonical points and shapes are still re-ordered into tile row groups -- that is what makes spatial subsetting cheap from Python -- but keep only their own columns. display_xy and feature_code go to /trx, display_geometry and cell_code to /cell_seg, with the same row-group layout. This fixes two separate problems with one change: 1. SpatialData.write() round-trip. A nested Arrow column cannot survive dask's parquet round-trip: it either fails outright on a schema mismatch or silently comes back as a string. A tiled store could therefore not be rewritten. Now it can, and a test asserts it. 2. Column projection. parquet-wasm corrupts the IPC stream whenever 'columns' is passed (0.7.1 and 0.7.2, apache-arrow 15 and 18, scalar and nested alike, even an empty array), which is what left the viewer showing only the image layer. A standalone render file makes projection unnecessary: reading every column of it already transfers only what is drawn. Also in this change: - All image channels are written, not just the first. Xenium morphology_focus has four; each gets its own pyramid and image_info entry with a distinct colour. Channel names are sanitised for paths, since names like 'ATP1A1/CD45/E-Cadherin' would otherwise create nested directories. - Render files are written before the canonical rewrite. Rewriting canonical replaces the files the lazy dask frame still points at, so the other order made the second read fail. - The streaming path reported render_only=False regardless; both paths now report it, with a test covering each. - Re-tiling drops render columns left in canonical by the previous layout. Co-Authored-By: Claude Opus 5 (1M context) --- .../experimental/points_parquet.py | 69 +++- .../experimental/shapes_parquet.py | 20 +- .../experimental/tiled_access.py | 136 +++++--- tests/test_points_parquet.py | 294 ++++++++++-------- tests/test_shapes_parquet.py | 99 +++++- tests/test_tiled_access.py | 99 +++++- 6 files changed, 510 insertions(+), 207 deletions(-) diff --git a/src/spatialdata_io/experimental/points_parquet.py b/src/spatialdata_io/experimental/points_parquet.py index 2d756ae4..152a61b5 100644 --- a/src/spatialdata_io/experimental/points_parquet.py +++ b/src/spatialdata_io/experimental/points_parquet.py @@ -141,8 +141,13 @@ def _prepare_table( feature_key: str, grid: RegularGrid, categories: Any | None, + render_only: bool = False, ) -> tuple[pa.Table, NDArray[np.int64]]: - """Add the render columns to one chunk and return it with its tile assignment. + """Build one chunk's output table and return it with its tile assignment. + + With ``render_only`` the table holds just the render columns, for the standalone file + a viewer reads. Otherwise it is the canonical columns, tile-ordered but otherwise + untouched. ``categories`` pins the categorical dictionary so that every chunk converts to an identical Arrow schema; without it, partitions observing different feature subsets @@ -150,11 +155,27 @@ def _prepare_table( """ px, py = _to_display_pixels(df["x"].to_numpy(), df["y"].to_numpy(), transform) tile_ids = grid.assign(px, py) - codes = catalog.encode(df[feature_key]) - # Re-tiling an already-tiled element must replace the render columns, not append - # duplicates: a duplicated name makes projected reads fail ("Multiple matches for - # FieldRef"), and a pandas round-trip degrades fixed_size_list to a variable list. + if render_only: + # Only the render columns. A viewer reads every column of this file, which is + # why no column projection is needed -- and parquet-wasm's projection is broken + # anyway (any `columns` argument corrupts the IPC stream it emits). + table = pa.table( + { + POSITION_COLUMN: _interleaved_positions(px, py), + FEATURE_COLUMN: pa.array(catalog.encode(df[feature_key])), + } + ) + return table, tile_ids + + # The canonical element keeps only its own columns. The render columns live in a + # separate file, for two reasons: a nested Arrow column cannot survive dask's parquet + # round-trip (SpatialData.write() either fails or silently returns it as a string), + # and a standalone render file means a viewer reads every column of it, so no column + # projection is needed -- which matters because parquet-wasm's projection is broken. + # + # Any render columns left by an earlier version are dropped, so re-tiling a store + # written before this change cleans it up rather than preserving them. stale = [c for c in (POSITION_COLUMN, FEATURE_COLUMN, _TILE_ID) if c in df.columns] if df.attrs or stale or categories is not None: df = df.copy(deep=False) @@ -166,10 +187,7 @@ def _prepare_table( if categories is not None and isinstance(df[feature_key].dtype, pd.CategoricalDtype): df[feature_key] = df[feature_key].cat.set_categories(categories) - table = pa.Table.from_pandas(df, preserve_index=True) - table = table.append_column(POSITION_COLUMN, _interleaved_positions(px, py)) - table = table.append_column(FEATURE_COLUMN, pa.array(codes)) - return table, tile_ids + return pa.Table.from_pandas(df, preserve_index=True), tile_ids def _sorted_by_tile(table: pa.Table, tile_ids: NDArray[np.int64]) -> tuple[pa.Table, NDArray[np.int64]]: @@ -229,6 +247,7 @@ def write_points_regular_grid( max_row_groups_per_file: int = DEFAULT_MAX_ROW_GROUPS_PER_FILE, compression: str = "zstd", streaming: bool | None = None, + render_only: bool = False, overwrite: bool = False, ) -> dict[str, Any]: """Write a Points element as regular-grid row groups. @@ -298,6 +317,7 @@ def write_points_regular_grid( feature_key=feature_key, max_row_groups_per_file=max_row_groups_per_file, compression=compression, + render_only=render_only, ) df = points.compute() if hasattr(points, "compute") else points @@ -305,7 +325,13 @@ def write_points_regular_grid( raise TypeError(f"expected a DataFrame, got {type(df).__name__}") table, tile_ids = _prepare_table( - df, transform=transform, catalog=catalog, feature_key=feature_key, grid=grid, categories=None + df, + transform=transform, + catalog=catalog, + feature_key=feature_key, + grid=grid, + categories=None, + render_only=render_only, ) table, sorted_tile_ids = _sorted_by_tile(table, tile_ids) @@ -348,7 +374,14 @@ def write_points_regular_grid( staging.rename(output_dir) return _manifest_fragment( - output_dir, filenames, grid, transform, catalog, max_row_groups_per_file, int(table.num_rows) + output_dir, + filenames, + grid, + transform, + catalog, + max_row_groups_per_file, + int(table.num_rows), + render_only=render_only, ) @@ -370,8 +403,9 @@ def _manifest_fragment( catalog: FeatureCatalog, max_row_groups_per_file: int, n_rows: int, + render_only: bool = False, ) -> dict[str, Any]: - return { + fragment: dict[str, Any] = { "directory": str(output_dir.name), "files": filenames, "max_row_groups_per_file": max_row_groups_per_file, @@ -381,14 +415,13 @@ def _manifest_fragment( "position_dtype": "uint32", "position_size": 2, "feature_column": FEATURE_COLUMN, - # Projected by the client, so canonical coordinates, ids and QC columns are never - # decoded or transferred during ordinary rendering. - "columns": [POSITION_COLUMN, FEATURE_COLUMN], "n_rows": n_rows, "tile_grid": grid.to_manifest_dict(), "display_transform": transform.to_manifest_dict(), "feature_catalog": catalog.to_manifest_dict(), + "render_only": render_only, } + return fragment def _write_streaming( @@ -401,6 +434,7 @@ def _write_streaming( feature_key: str, max_row_groups_per_file: int, compression: str, + render_only: bool = False, ) -> dict[str, Any]: """Write the tiled output without holding the whole element in memory. @@ -439,6 +473,7 @@ def _write_streaming( feature_key=feature_key, grid=grid, categories=categories, + render_only=render_only, ) n_rows += table.num_rows if schema is None: @@ -501,4 +536,6 @@ def _write_streaming( shutil.rmtree(output_dir) staging.rename(output_dir) - return _manifest_fragment(output_dir, filenames, grid, transform, catalog, max_row_groups_per_file, n_rows) + return _manifest_fragment( + output_dir, filenames, grid, transform, catalog, max_row_groups_per_file, n_rows, render_only=render_only + ) diff --git a/src/spatialdata_io/experimental/shapes_parquet.py b/src/spatialdata_io/experimental/shapes_parquet.py index 654ea4a1..fa792b89 100644 --- a/src/spatialdata_io/experimental/shapes_parquet.py +++ b/src/spatialdata_io/experimental/shapes_parquet.py @@ -130,6 +130,7 @@ def write_shapes_regular_grid( cell_index: Any | None = None, max_row_groups_per_file: int = DEFAULT_MAX_ROW_GROUPS_PER_FILE, compression: str = "zstd", + render_only: bool = False, overwrite: bool = False, ) -> dict[str, Any]: """Write a Shapes element as regular-grid row groups. @@ -168,11 +169,11 @@ def write_shapes_regular_grid( transform = display_transform or DisplayTransform.from_element(shapes, coordinate_system) display, cx, cy = _display_geometry_array(shapes.geometry, transform) - # Re-tiling an already-tiled element must replace the render columns, not append - # duplicates: a duplicated name makes projected reads fail outright, and the stale - # copy is also mistyped because a pandas round-trip degrades fixed_size_list to list. + # Any render columns left by an earlier version are dropped, so re-tiling a store + # written before the split cleans it up instead of preserving them. stale = [c for c in (GEOMETRY_COLUMN, CELL_CODE_COLUMN) if c in shapes.columns] - table = _canonical_geoparquet_table(shapes.drop(columns=stale) if stale else shapes) + canonical = shapes.drop(columns=stale) if stale else shapes + table = None if render_only else _canonical_geoparquet_table(canonical) if cell_index is None: codes = np.arange(len(shapes), dtype=np.uint32) @@ -186,8 +187,12 @@ def write_shapes_regular_grid( ) codes = np.fromiter((positions[k] for k in shapes.index), dtype=np.uint32, count=len(shapes)) - table = table.append_column(GEOMETRY_COLUMN, display) - table = table.append_column(CELL_CODE_COLUMN, pa.array(codes)) + if render_only: + # Standalone render file: a viewer reads every column, so no projection is needed + # and the canonical GeoParquet keeps only its own WKB geometry. + table = pa.table({GEOMETRY_COLUMN: display, CELL_CODE_COLUMN: pa.array(codes)}) + else: + table = table.append_column(CELL_CODE_COLUMN, pa.array(codes)) tile_ids = grid.assign(cx, cy) order = np.argsort(tile_ids, kind="stable") @@ -246,8 +251,7 @@ def write_shapes_regular_grid( fragment: dict[str, Any] = { "geometry_column": GEOMETRY_COLUMN, "cell_id_column": CELL_CODE_COLUMN, - # Projected by the client: the canonical WKB geometry is never transferred. - "columns": [GEOMETRY_COLUMN, CELL_CODE_COLUMN], + "render_only": render_only, "max_row_groups_per_file": max_row_groups_per_file, "total_row_groups": grid.num_tiles, "n_shapes": int(table.num_rows), diff --git a/src/spatialdata_io/experimental/tiled_access.py b/src/spatialdata_io/experimental/tiled_access.py index 4c658cff..638ea831 100644 --- a/src/spatialdata_io/experimental/tiled_access.py +++ b/src/spatialdata_io/experimental/tiled_access.py @@ -50,6 +50,42 @@ #: Directory inside the store holding derived (non-canonical) profile assets. PROFILE_DIR = "visualization" +#: Display colours cycled through when a channel has none assigned. First is blue, which +#: is the conventional nuclear stain colour and usually channel 0 (DAPI). +_DEFAULT_CHANNEL_COLORS = [ + (0, 0, 255), + (0, 255, 0), + (255, 0, 0), + (255, 255, 0), + (255, 0, 255), + (0, 255, 255), +] + + +def _channels_of(element: Any) -> list[Any]: + """List an image element's channel names, falling back to indices.""" + level = element[next(iter(element.children))] if hasattr(element, "children") else element + array = level[next(iter(level.data_vars))] if hasattr(level, "data_vars") else level + coords = getattr(array, "coords", {}) + if "c" in coords: + return [str(c) for c in coords["c"].values] + size = dict(zip(array.dims, array.shape, strict=True)).get("c", 1) + return list(range(size)) + + +def _channel_label(channel: Any, index: int) -> str: + """A filesystem- and URL-safe label for a channel. + + Xenium channel names include slashes ('ATP1A1/CD45/E-Cadherin'), which would other- + wise create nested directories and break the manifest's relative paths. + """ + if isinstance(channel, int): + return f"channel_{channel}" + safe = "".join(ch if ch.isalnum() else "_" for ch in str(channel)).strip("_").lower() + while "__" in safe: + safe = safe.replace("__", "_") + return safe or f"channel_{index}" + def _grid_for(points: Any, transform: DisplayTransform, tile_size_px: float) -> RegularGrid: """Derive the grid covering the points element in display pixel space.""" @@ -72,10 +108,8 @@ def add_spatial_tiling( technology: str = "Xenium", include_cbg: bool = True, image_element: str | None = None, - image_channel: int | str | None = None, - image_name: str = "dapi", - image_button_name: str = "DAPI", - image_color: tuple[int, int, int] = (0, 0, 255), + image_channels: list[int | str] | None = None, + image_colors: dict[str, tuple[int, int, int]] | None = None, image_tile_size: int = 512, compression: str = "zstd", ) -> dict[str, Any]: @@ -111,10 +145,11 @@ def add_spatial_tiling( image_element Name of an image element to render into a WebP display pyramid, or ``None`` to skip images. The canonical OME-Zarr image is left untouched either way. - image_channel - Channel index or name to render. Defaults to the first channel. - image_name, image_button_name, image_color - Celldega channel descriptor recorded in ``image_info``. + image_channels + Channels to render, by name or index. Defaults to every channel in the element. + image_colors + Optional display colour per channel name. Channels without an entry get a colour + from a default palette. image_tile_size Image tile edge length in pixels. compression @@ -147,7 +182,24 @@ def add_spatial_tiling( profile_dir = store / PROFILE_DIR / PROFILE_NAME profile_dir.mkdir(parents=True, exist_ok=True) + # The render columns go to a standalone file inside the profile directory. A viewer + # reads every column of it, so it needs no column projection, and the canonical + # element is left free of nested Arrow columns. transcripts = write_points_regular_grid( + points, + profile_dir / "trx", + catalog=catalog, + grid=grid, + display_transform=transform, + feature_key=feature_key, + max_row_groups_per_file=max_row_groups_per_file, + compression=compression, + render_only=True, + overwrite=True, + ) + # The canonical element is re-ordered into tile row groups but keeps only its own + # columns, so it still round-trips through SpatialData.write() and normal reads. + write_points_regular_grid( points, store / "points" / points_element / "points.parquet", catalog=catalog, @@ -158,9 +210,6 @@ def add_spatial_tiling( compression=compression, overwrite=True, ) - # Paths in the manifest are relative to the profile directory, so Celldega can be - # pointed at that directory as its base_url with no reader change. - transcripts["directory"] = f"../../points/{points_element}/points.parquet" cell_segmentation = None cell_metadata = None @@ -182,6 +231,17 @@ def add_spatial_tiling( ) cell_names = [str(k) for k in (table.obs_names if table is not None else shapes.index)] cell_segmentation = write_shapes_regular_grid( + shapes, + profile_dir / "cell_seg", + grid=grid, + display_transform=shapes_transform, + cell_index=list(table.obs_names) if table is not None else None, + max_row_groups_per_file=max_row_groups_per_file, + compression=compression, + render_only=True, + overwrite=True, + ) + write_shapes_regular_grid( shapes, store / "shapes" / shapes_element / "shapes.parquet", grid=grid, @@ -191,11 +251,6 @@ def add_spatial_tiling( compression=compression, overwrite=True, ) - prefix = f"../../shapes/{shapes_element}" - if "path" in cell_segmentation: - cell_segmentation["path"] = f"{prefix}/{cell_segmentation['path']}" - else: - cell_segmentation["directory"] = f"{prefix}/{cell_segmentation['directory']}" cbg = None if include_cbg and table is not None: @@ -217,25 +272,36 @@ def add_spatial_tiling( if image_element not in sdata.images: raise ValueError(f"image element {image_element!r} not found; have {list(sdata.images)}") - pyramid = write_webp_pyramid( - sdata.images[image_element], - profile_dir / "images" / image_name, - channel=image_channel, - tile_size=image_tile_size, - source_element=image_element, - overwrite=True, - ) - # ImageRowGroupReader resolves files as baseUrl/directory/file and reads the - # per-zoom grid from the entry's zoom_info. - pyramid["directory"] = f"images/{image_name}" - images[image_name] = pyramid - image_info = [{"name": image_name, "button_name": image_button_name, "color": list(image_color)}] - image_dimensions = { - "width": pyramid["source_width"], - "height": pyramid["source_height"], - "tile_size": image_tile_size, - } - max_pyramid_zoom = pyramid["max_zoom"] + element = sdata.images[image_element] + channels = image_channels if image_channels is not None else _channels_of(element) + + for index, channel in enumerate(channels): + label = _channel_label(channel, index) + pyramid = write_webp_pyramid( + element, + profile_dir / "images" / label, + channel=channel, + tile_size=image_tile_size, + source_element=image_element, + overwrite=True, + ) + # ImageRowGroupReader resolves files as baseUrl/directory/file and reads the + # per-zoom grid from the entry's zoom_info. + pyramid["directory"] = f"images/{label}" + images[label] = pyramid + colour = (image_colors or {}).get(label) or _DEFAULT_CHANNEL_COLORS[ + index % len(_DEFAULT_CHANNEL_COLORS) + ] + image_info.append( + {"name": label, "button_name": str(channel), "color": list(colour)} + ) + # Every channel of one element shares its dimensions and pyramid depth. + image_dimensions = { + "width": pyramid["source_width"], + "height": pyramid["source_height"], + "tile_size": image_tile_size, + } + max_pyramid_zoom = pyramid["max_zoom"] catalog.to_frame().to_parquet(profile_dir / "meta_gene.parquet", index=False) diff --git a/tests/test_points_parquet.py b/tests/test_points_parquet.py index 278944a6..1a9a5b01 100644 --- a/tests/test_points_parquet.py +++ b/tests/test_points_parquet.py @@ -1,8 +1,16 @@ """Tests for the regular-grid Points rewrite. -The synthetic fixture is a 2x3 tile grid deliberately containing the awkward cases: -an empty tile, a point exactly on a tile boundary, several points sharing one tile, -and both genes and a non-gene control feature. +Two outputs are produced from one element and both are tested here: + +* the **canonical** file, re-ordered into tile row groups but carrying only its own + columns, so it still round-trips through ``SpatialData.write()``; +* the **render** file, holding only ``display_xy`` and ``feature_code``, which a viewer + reads in full -- no column projection, because parquet-wasm's projection corrupts the + IPC stream it emits. + +The synthetic fixture is a 2x3 tile grid deliberately containing the awkward cases: an +empty tile, a point exactly on a tile boundary, several points sharing one tile, and both +genes and a non-gene control feature. """ from __future__ import annotations @@ -48,7 +56,7 @@ def catalog() -> FeatureCatalog: @pytest.fixture def points() -> pd.DataFrame: - """A pandas Points-like frame with canonical micron coords and extra annotations.""" + """A Points element with canonical micron coords and extra annotations.""" px = np.array([s[0] for s in POINTS_SPEC]) py = np.array([s[1] for s in POINTS_SPEC]) df = pd.DataFrame( @@ -63,66 +71,84 @@ def points() -> pd.DataFrame: } ) from spatialdata.models import PointsModel + from spatialdata.transformations import Scale, set_transformation - return PointsModel.parse(df, coordinates={"x": "x", "y": "y", "z": "z"}, feature_key="feature_name") + element = PointsModel.parse(df, coordinates={"x": "x", "y": "y", "z": "z"}, feature_key="feature_name") + set_transformation(element, Scale([2.0, 2.0], axes=("x", "y")), "global") + return element @pytest.fixture def written(tmp_path: Path, points: pd.DataFrame, catalog: FeatureCatalog) -> tuple[Path, dict]: - from spatialdata.transformations import Scale, set_transformation - - set_transformation(points, Scale([2.0, 2.0], axes=("x", "y")), "global") + """The canonical element: tile-ordered, carrying only its own columns.""" out = tmp_path / "points.parquet" - manifest = write_points_regular_grid(points, out, catalog=catalog, grid=GRID) - return out, manifest + return out, write_points_regular_grid(points, out, catalog=catalog, grid=GRID) + + +@pytest.fixture +def rendered(tmp_path: Path, points: pd.DataFrame, catalog: FeatureCatalog) -> tuple[Path, dict]: + """The standalone render file: display_xy and feature_code only.""" + out = tmp_path / "trx" + return out, write_points_regular_grid(points, out, catalog=catalog, grid=GRID, render_only=True) def _read_all(directory: Path, manifest: dict) -> pa.Table: return pa.concat_tables([pq.read_table(directory / f) for f in manifest["files"]]) +def _row_group(directory: Path, manifest: dict, tile_id: int) -> pa.Table: + file_index, local = GRID.chunk_location(tile_id, manifest["max_row_groups_per_file"]) + return pq.ParquetFile(directory / manifest["files"][file_index]).read_row_group(local) + + # -- row-group layout --------------------------------------------------------- -def test_row_group_count_equals_tile_count(written: tuple[Path, dict]) -> None: - directory, manifest = written +@pytest.mark.parametrize("fixture", ["written", "rendered"]) +def test_row_group_count_equals_tile_count(fixture: str, request: pytest.FixtureRequest) -> None: + directory, manifest = request.getfixturevalue(fixture) total = sum(pq.ParquetFile(directory / f).metadata.num_row_groups for f in manifest["files"]) assert total == GRID.num_tiles == 6 assert manifest["total_row_groups"] == 6 -def test_each_tile_row_group_holds_exactly_its_points(written: tuple[Path, dict]) -> None: +@pytest.mark.parametrize("fixture", ["written", "rendered"]) +def test_each_tile_row_group_holds_exactly_its_points(fixture: str, request: pytest.FixtureRequest) -> None: """The core contract: row_group_index == tile_id, with no lookup table.""" - directory, manifest = written - expected: dict[int, list[tuple[float, float]]] = {} - for pxv, pyv, _, tid in POINTS_SPEC: - expected.setdefault(tid, []).append((pxv, pyv)) - + directory, manifest = request.getfixturevalue(fixture) + expected: dict[int, int] = {} + for *_, tid in POINTS_SPEC: + expected[tid] = expected.get(tid, 0) + 1 for tile_id in range(GRID.num_tiles): - file_index, local = GRID.chunk_location(tile_id, manifest["max_row_groups_per_file"]) - f = pq.ParquetFile(directory / manifest["files"][file_index]) - rg = f.read_row_group(local, columns=[POSITION_COLUMN]) - got = [tuple(v) for v in rg[POSITION_COLUMN].to_pylist()] - assert sorted(got) == sorted(expected.get(tile_id, [])), f"tile {tile_id}" + got = _row_group(directory, manifest, tile_id).num_rows + assert got == expected.get(tile_id, 0), f"tile {tile_id}" -def test_empty_tile_is_a_zero_row_row_group(written: tuple[Path, dict]) -> None: - directory, manifest = written +@pytest.mark.parametrize("fixture", ["written", "rendered"]) +def test_empty_tile_is_a_zero_row_row_group(fixture: str, request: pytest.FixtureRequest) -> None: + directory, manifest = request.getfixturevalue(fixture) file_index, local = GRID.chunk_location(4, manifest["max_row_groups_per_file"]) - f = pq.ParquetFile(directory / manifest["files"][file_index]) - assert f.metadata.row_group(local).num_rows == 0 - assert f.read_row_group(local).num_rows == 0 + md = pq.ParquetFile(directory / manifest["files"][file_index]).metadata + assert md.row_group(local).num_rows == 0 -def test_boundary_point_goes_to_the_upper_tile(written: tuple[Path, dict]) -> None: +def test_boundary_point_goes_to_the_upper_tile(rendered: tuple[Path, dict]) -> None: """A point at exactly x=10 belongs to tile_x=1 under half-open bounds.""" - directory, manifest = written - f = pq.ParquetFile(directory / manifest["files"][0]) - rg = f.read_row_group(3, columns=[POSITION_COLUMN]) + directory, manifest = rendered + rg = _row_group(directory, manifest, 3) assert [tuple(v) for v in rg[POSITION_COLUMN].to_pylist()] == [(10, 2)] -# -- canonical data preservation ---------------------------------------------- +# -- the canonical element ---------------------------------------------------- + + +def test_canonical_file_has_no_render_columns(written: tuple[Path, dict]) -> None: + """A nested Arrow column cannot survive dask's parquet round-trip, so it stays out.""" + directory, manifest = written + names = _read_all(directory, manifest).column_names + assert POSITION_COLUMN not in names + assert FEATURE_COLUMN not in names + assert manifest["render_only"] is False def test_no_row_lost_or_duplicated(written: tuple[Path, dict], points: pd.DataFrame) -> None: @@ -141,7 +167,7 @@ def test_canonical_columns_are_unchanged(written: tuple[Path, dict], points: pd. merged = got.set_index("transcript_id").loc[original["transcript_id"].to_numpy()] for col in ("x", "y", "z", "cell_id", "qv"): - np.testing.assert_array_equal(merged[col].to_numpy(), original[col].to_numpy(), err_msg=f"column {col} changed") + np.testing.assert_array_equal(merged[col].to_numpy(), original[col].to_numpy(), err_msg=col) assert list(merged["feature_name"].astype(str)) == list(original["feature_name"].astype(str)) @@ -152,29 +178,34 @@ def test_index_is_preserved(written: tuple[Path, dict], points: pd.DataFrame) -> assert sorted(got.index.tolist()) == sorted(original.index.tolist()) -def test_row_order_is_grouped_by_tile(written: tuple[Path, dict]) -> None: +def test_canonical_rows_are_grouped_by_tile(written: tuple[Path, dict]) -> None: + """The reordering is the point: it is what makes spatial subsetting cheap in Python.""" directory, manifest = written - table = _read_all(directory, manifest) - xs = np.array([v[0] for v in table[POSITION_COLUMN].to_pylist()]) - ys = np.array([v[1] for v in table[POSITION_COLUMN].to_pylist()]) - tile_ids = GRID.assign(xs, ys) - assert (np.diff(tile_ids) >= 0).all(), "rows are not grouped by tile" + df = _read_all(directory, manifest).to_pandas() + tile_ids = GRID.assign(df["x"].to_numpy() * 2, df["y"].to_numpy() * 2) + assert (np.diff(tile_ids) >= 0).all() -# -- render columns ----------------------------------------------------------- +# -- the render file ---------------------------------------------------------- -def test_display_xy_is_fixed_size_list_uint32(written: tuple[Path, dict]) -> None: - directory, manifest = written +def test_render_file_holds_only_the_render_columns(rendered: tuple[Path, dict]) -> None: + """A viewer reads every column of this file, which is why no projection is needed.""" + directory, manifest = rendered + assert _read_all(directory, manifest).column_names == [POSITION_COLUMN, FEATURE_COLUMN] + assert manifest["render_only"] is True + + +def test_display_xy_is_fixed_size_list_uint32(rendered: tuple[Path, dict]) -> None: + directory, manifest = rendered field = pq.ParquetFile(directory / manifest["files"][0]).schema_arrow.field(POSITION_COLUMN) assert field.type == pa.list_(pa.uint32(), 2) assert manifest["position_dtype"] == "uint32" - assert manifest["position_encoding"] == "fixed_size_list" -def test_display_xy_child_buffer_is_interleaved(written: tuple[Path, dict]) -> None: +def test_display_xy_child_buffer_is_interleaved(rendered: tuple[Path, dict]) -> None: """The flat child buffer must be [x0,y0,x1,y1,...] so deck.gl can consume it directly.""" - directory, manifest = written + directory, manifest = rendered col = _read_all(directory, manifest)[POSITION_COLUMN].combine_chunks() flat = col.values.to_numpy(zero_copy_only=False) pairs = [tuple(v) for v in col.to_pylist()] @@ -182,51 +213,41 @@ def test_display_xy_child_buffer_is_interleaved(written: tuple[Path, dict]) -> N assert flat[1::2].tolist() == [p[1] for p in pairs] -def test_display_xy_matches_the_transform(written: tuple[Path, dict]) -> None: - directory, manifest = written - table = _read_all(directory, manifest) - got = {int(t): tuple(v) for t, v in zip(table["transcript_id"].to_pylist(), table[POSITION_COLUMN].to_pylist())} - for i, (pxv, pyv, _, _) in enumerate(POINTS_SPEC): - assert got[100 + i] == (int(pxv), int(pyv)) +def test_display_xy_matches_the_transform(rendered: tuple[Path, dict]) -> None: + directory, manifest = rendered + got = sorted(tuple(v) for v in _read_all(directory, manifest)[POSITION_COLUMN].to_pylist()) + assert got == sorted((int(x), int(y)) for x, y, _, _ in POINTS_SPEC) -def test_feature_codes_match_catalog(written: tuple[Path, dict], catalog: FeatureCatalog) -> None: - directory, manifest = written +def _codes_by_position(directory: Path, manifest: dict) -> dict[tuple[int, int], int]: table = _read_all(directory, manifest) - for name, code in zip(table["feature_name"].to_pylist(), table[FEATURE_COLUMN].to_pylist()): - assert catalog.names[code] == name + return {tuple(v): c for v, c in zip(table[POSITION_COLUMN].to_pylist(), table[FEATURE_COLUMN].to_pylist())} -def test_control_feature_is_coded_above_every_gene(written: tuple[Path, dict], catalog: FeatureCatalog) -> None: - directory, manifest = written - table = _read_all(directory, manifest).to_pandas() - control = table[table["feature_name"].astype(str).str.startswith("NegControl")] - assert len(control) == 1 - assert int(control[FEATURE_COLUMN].iloc[0]) >= catalog.n_genes +def test_feature_codes_match_catalog(rendered: tuple[Path, dict], catalog: FeatureCatalog) -> None: + codes = _codes_by_position(*rendered) + for x, y, feature, _ in POINTS_SPEC: + assert catalog.names[codes[(int(x), int(y))]] == feature -# -- column projection -------------------------------------------------------- +def test_control_feature_is_coded_above_every_gene(rendered: tuple[Path, dict], catalog: FeatureCatalog) -> None: + codes = _codes_by_position(*rendered) + assert codes[(15, 28)] >= catalog.n_genes -def test_column_projection_reads_only_render_columns(written: tuple[Path, dict]) -> None: - """Celldega projects these two columns; canonical columns must not be required.""" - directory, manifest = written - f = pq.ParquetFile(directory / manifest["files"][0]) - rg = f.read_row_group(0, columns=[POSITION_COLUMN, FEATURE_COLUMN]) - assert rg.column_names == [POSITION_COLUMN, FEATURE_COLUMN] - assert rg.num_rows == 3 +def test_render_file_is_smaller_than_canonical(written: tuple[Path, dict], rendered: tuple[Path, dict]) -> None: + """The render file is what crosses the wire, so it should be a fraction of canonical.""" + canonical = sum(f.stat().st_size for f in written[0].glob("*.parquet")) + render = sum(f.stat().st_size for f in rendered[0].glob("*.parquet")) + assert render < canonical # -- file layout -------------------------------------------------------------- def test_multi_file_split_and_padding(tmp_path: Path, points: pd.DataFrame, catalog: FeatureCatalog) -> None: - from spatialdata.transformations import Scale, set_transformation - - set_transformation(points, Scale([2.0, 2.0], axes=("x", "y")), "global") out = tmp_path / "multi.parquet" manifest = write_points_regular_grid(points, out, catalog=catalog, grid=GRID, max_row_groups_per_file=2) - assert manifest["files"] == ["chunk_0.parquet", "chunk_1.parquet", "chunk_2.parquet"] assert [pq.ParquetFile(out / f).metadata.num_row_groups for f in manifest["files"]] == [2, 2, 2] assert sum(pq.ParquetFile(out / f).metadata.num_rows for f in manifest["files"]) == len(POINTS_SPEC) @@ -247,13 +268,10 @@ def test_overwrite_guard(written: tuple[Path, dict], points: pd.DataFrame, catal def test_failure_leaves_no_partial_output(tmp_path: Path, points: pd.DataFrame) -> None: """A mid-write error must not leave a half-rewritten directory behind.""" - from spatialdata.transformations import Scale, set_transformation - - set_transformation(points, Scale([2.0, 2.0], axes=("x", "y")), "global") out = tmp_path / "boom.parquet" bad = FeatureCatalog(names=("GENEA",), n_genes=1) # missing GENEB -> encode() raises with pytest.raises(ValueError, match="not in the catalog"): - write_points_regular_grid(points, out, catalog=bad, grid=GRID) + write_points_regular_grid(points, out, catalog=bad, grid=GRID, render_only=True) assert not out.exists() assert not out.with_name(out.name + ".tmp").exists() @@ -277,86 +295,58 @@ def test_missing_feature_column_is_reported(tmp_path: Path, points: pd.DataFrame write_points_regular_grid(points, tmp_path / "x.parquet", catalog=catalog, grid=GRID, feature_key="nope") -def test_rewrite_is_idempotent(tmp_path: Path, points: pd.DataFrame, catalog: FeatureCatalog) -> None: - """Re-optimizing an already-optimized element replaces the render columns, not appends them. - - A duplicated column name makes projected reads fail outright, and pandas round-trips - fixed_size_list back as a variable-length list, so a stale copy is also mistyped. - """ - from spatialdata.transformations import Scale, set_transformation +# -- streaming ---------------------------------------------------------------- - set_transformation(points, Scale([2.0, 2.0], axes=("x", "y")), "global") - out = tmp_path / "idem.parquet" - m1 = write_points_regular_grid(points, out, catalog=catalog, grid=GRID) - first = _read_all(out, m1) - # Feed the written result back in, the way read_zarr would hand it back: a dask - # frame that already carries display_xy/feature_code. +def _as_partitioned(frame: pd.DataFrame, npartitions: int = 3): import dask.dataframe as dd + from spatialdata.transformations import Scale - again = dd.from_pandas(first.to_pandas(), npartitions=1) - # Attach the transform the way read_zarr does, rather than via set_transformation, - # which requires an element that already carries one. - again.attrs["transform"] = {"global": Scale([2.0, 2.0], axes=("x", "y"))} - assert POSITION_COLUMN in again.columns # precondition: the stale columns are present - m2 = write_points_regular_grid(again, out, catalog=catalog, grid=GRID, overwrite=True) - second = _read_all(out, m2) - - assert second.column_names.count(POSITION_COLUMN) == 1 - assert second.column_names.count(FEATURE_COLUMN) == 1 - assert second.schema.field(POSITION_COLUMN).type == pa.list_(pa.uint32(), 2) - assert second.num_rows == first.num_rows - assert second[POSITION_COLUMN].to_pylist() == first[POSITION_COLUMN].to_pylist() + pdf = frame.compute() if hasattr(frame, "compute") else frame + out = dd.from_pandas(pdf, npartitions=npartitions) + out.attrs["transform"] = {"global": Scale([2.0, 2.0], axes=("x", "y"))} + return out -def test_streaming_matches_in_memory(tmp_path: Path, points: pd.DataFrame, catalog: FeatureCatalog) -> None: +@pytest.mark.parametrize("render_only", [False, True]) +def test_streaming_matches_in_memory( + tmp_path: Path, points: pd.DataFrame, catalog: FeatureCatalog, render_only: bool +) -> None: """The two write paths must be interchangeable, or large datasets would diverge. - Grouping by tile is a global sort, so the streaming path spills rows into per-output - -file buckets and sorts each independently. That must land every row in the same row - group as the single-pass path. + Grouping by tile is a global sort, so the streaming path spills rows into + per-output-file buckets and sorts each independently. That must land every row in the + same row group as the single-pass path. """ - import dask.dataframe as dd - from spatialdata.transformations import Scale, set_transformation - - set_transformation(points, Scale([2.0, 2.0], axes=("x", "y")), "global") - in_memory = tmp_path / "mem.parquet" - m1 = write_points_regular_grid(points, in_memory, catalog=catalog, grid=GRID, streaming=False) - - # Several partitions, so the spill path is genuinely exercised. - chunked = dd.from_pandas(points.compute(), npartitions=3) - chunked.attrs["transform"] = {"global": Scale([2.0, 2.0], axes=("x", "y"))} - streamed = tmp_path / "stream.parquet" + in_memory = tmp_path / "mem" + m1 = write_points_regular_grid( + points, in_memory, catalog=catalog, grid=GRID, streaming=False, render_only=render_only + ) + streamed = tmp_path / "stream" m2 = write_points_regular_grid( - chunked, streamed, catalog=catalog, grid=GRID, streaming=True, max_row_groups_per_file=2 + _as_partitioned(points), + streamed, + catalog=catalog, + grid=GRID, + streaming=True, + max_row_groups_per_file=2, + render_only=render_only, ) assert m1["total_row_groups"] == m2["total_row_groups"] == GRID.num_tiles assert m1["n_rows"] == m2["n_rows"] == len(POINTS_SPEC) - - # Every tile must hold exactly the same points in both. for tile_id in range(GRID.num_tiles): - fi1, lo1 = GRID.chunk_location(tile_id, m1["max_row_groups_per_file"]) - fi2, lo2 = GRID.chunk_location(tile_id, m2["max_row_groups_per_file"]) - a = pq.ParquetFile(in_memory / m1["files"][fi1]).read_row_group(lo1, columns=[POSITION_COLUMN]) - b = pq.ParquetFile(streamed / m2["files"][fi2]).read_row_group(lo2, columns=[POSITION_COLUMN]) - assert a[POSITION_COLUMN].to_pylist() == b[POSITION_COLUMN].to_pylist(), f"tile {tile_id}" + assert _row_group(in_memory, m1, tile_id).num_rows == _row_group(streamed, m2, tile_id).num_rows def test_streaming_preserves_canonical_columns(tmp_path: Path, points: pd.DataFrame, catalog: FeatureCatalog) -> None: - import dask.dataframe as dd - from spatialdata.transformations import Scale - - original = points.compute() - chunked = dd.from_pandas(original, npartitions=3) - chunked.attrs["transform"] = {"global": Scale([2.0, 2.0], axes=("x", "y"))} + original = points.compute() if hasattr(points, "compute") else points out = tmp_path / "s.parquet" - manifest = write_points_regular_grid(chunked, out, catalog=catalog, grid=GRID, streaming=True) + manifest = write_points_regular_grid(_as_partitioned(points), out, catalog=catalog, grid=GRID, streaming=True) got = _read_all(out, manifest).to_pandas().set_index("transcript_id").loc[original["transcript_id"].to_numpy()] for col in ("x", "y", "z", "cell_id", "qv"): np.testing.assert_array_equal(got[col].to_numpy(), original[col].to_numpy(), err_msg=col) - assert list(got["feature_name"].astype(str)) == list(original["feature_name"].astype(str)) def test_streaming_requires_a_partitioned_element( @@ -364,3 +354,37 @@ def test_streaming_requires_a_partitioned_element( ) -> None: with pytest.raises(ValueError, match="streaming requires a partitioned"): write_points_regular_grid(points.compute(), tmp_path / "x.parquet", catalog=catalog, grid=GRID, streaming=True) + + +def test_retiling_drops_render_columns_left_by_an_older_writer( + tmp_path: Path, points: pd.DataFrame, catalog: FeatureCatalog +) -> None: + """Stores written before the split carry render columns in the canonical element.""" + out = tmp_path / "idem.parquet" + first = write_points_regular_grid(points, out, catalog=catalog, grid=GRID) + stale = _read_all(out, first).to_pandas() + stale[POSITION_COLUMN] = [[1, 2]] * len(stale) + stale[FEATURE_COLUMN] = np.zeros(len(stale), dtype=np.uint16) + + again = write_points_regular_grid( + _as_partitioned(stale, npartitions=1), out, catalog=catalog, grid=GRID, overwrite=True + ) + names = _read_all(out, again).column_names + assert POSITION_COLUMN not in names + assert FEATURE_COLUMN not in names + + +@pytest.mark.parametrize("render_only", [False, True]) +def test_manifest_reports_which_flavour_was_written( + tmp_path: Path, points: pd.DataFrame, catalog: FeatureCatalog, render_only: bool +) -> None: + """Both write paths must report it; the streaming path previously always said False.""" + in_memory = write_points_regular_grid( + points, tmp_path / "mem", catalog=catalog, grid=GRID, streaming=False, render_only=render_only + ) + streamed = write_points_regular_grid( + _as_partitioned(points), tmp_path / "str", catalog=catalog, grid=GRID, + streaming=True, render_only=render_only, + ) + assert in_memory["render_only"] is render_only + assert streamed["render_only"] is render_only diff --git a/tests/test_shapes_parquet.py b/tests/test_shapes_parquet.py index e30a8e20..c8c068b5 100644 --- a/tests/test_shapes_parquet.py +++ b/tests/test_shapes_parquet.py @@ -59,11 +59,20 @@ def shapes() -> gpd.GeoDataFrame: @pytest.fixture def written(tmp_path: Path, shapes: gpd.GeoDataFrame) -> tuple[Path, dict]: + """The canonical element: tile-ordered GeoParquet, no nested render column.""" out = tmp_path / "shapes.parquet" manifest = write_shapes_regular_grid(shapes, out, grid=GRID, display_transform=XFORM) return out, manifest +@pytest.fixture +def rendered(tmp_path: Path, shapes: gpd.GeoDataFrame) -> tuple[Path, dict]: + """The standalone render file: display_geometry and cell_code only.""" + out = tmp_path / "cell_seg.parquet" + manifest = write_shapes_regular_grid(shapes, out, grid=GRID, display_transform=XFORM, render_only=True) + return out, manifest + + # -- layout ------------------------------------------------------------------- @@ -136,18 +145,18 @@ def test_non_geometry_columns_survive(written: tuple[Path, dict], shapes: gpd.Ge # -- display geometry --------------------------------------------------------- -def test_display_geometry_has_the_nested_layout(written: tuple[Path, dict]) -> None: +def test_display_geometry_has_the_nested_layout(rendered: tuple[Path, dict]) -> None: """polygon -> rings -> interleaved uint32 pairs, as get_polygon_data.js walks it.""" - out, _ = written + out, _ = rendered t = pq.ParquetFile(out).schema_arrow.field(GEOMETRY_COLUMN).type assert pa.types.is_list(t) # polygon level assert pa.types.is_list(t.value_type) # ring level assert t.value_type.value_type == pa.list_(pa.uint32(), 2) # interleaved vertices -def test_display_geometry_offsets_resolve_like_the_js_reader(written: tuple[Path, dict]) -> None: +def test_display_geometry_offsets_resolve_like_the_js_reader(rendered: tuple[Path, dict]) -> None: """Mirror of getPolygonDataFromChunk: polygon offset -> ring offset -> coord index.""" - out, _ = written + out, _ = rendered col = pq.read_table(out)[GEOMETRY_COLUMN].combine_chunks() polygon_offsets = col.offsets.to_numpy() rings = col.values @@ -161,14 +170,14 @@ def test_display_geometry_offsets_resolve_like_the_js_reader(written: tuple[Path assert first.tolist() == col.to_pylist()[0][0][0] -def test_display_geometry_is_exterior_ring_only(written: tuple[Path, dict]) -> None: - out, _ = written +def test_display_geometry_is_exterior_ring_only(rendered: tuple[Path, dict]) -> None: + out, _ = rendered for polygon in pq.read_table(out)[GEOMETRY_COLUMN].to_pylist(): assert len(polygon) == 1, "expected exactly one ring per display polygon" -def test_display_vertices_match_the_transform(written: tuple[Path, dict]) -> None: - out, _ = written +def test_display_vertices_match_the_transform(rendered: tuple[Path, dict]) -> None: + out, _ = rendered table = pq.read_table(out) codes = table[CELL_CODE_COLUMN].to_pylist() geoms = table[GEOMETRY_COLUMN].to_pylist() @@ -185,7 +194,7 @@ def test_multipolygon_reduces_to_largest_part(tmp_path: Path) -> None: big, small = _square(2.5, 2.5, 2.0), _square(8.0, 13.0, 0.5) gdf = gpd.GeoDataFrame(geometry=[MultiPolygon([big, small])], index=["multi"]) out = tmp_path / "m.parquet" - write_shapes_regular_grid(ShapesModel.parse(gdf), out, grid=GRID, display_transform=XFORM) + write_shapes_regular_grid(ShapesModel.parse(gdf), out, grid=GRID, display_transform=XFORM, render_only=True) poly = pq.read_table(out)[GEOMETRY_COLUMN].to_pylist()[0] got = {tuple(v) for v in poly[0]} assert got == {(int(round(x * 2)), int(round(y * 2))) for x, y in big.exterior.coords} @@ -215,3 +224,75 @@ def test_overwrite_guard(written: tuple[Path, dict], shapes: gpd.GeoDataFrame) - out, _ = written with pytest.raises(FileExistsError): write_shapes_regular_grid(shapes, out, grid=GRID, display_transform=XFORM) + + +# -- cell metadata ------------------------------------------------------------ + + +def test_cell_metadata_schema_matches_celldega(tmp_path: Path, shapes: gpd.GeoDataFrame) -> None: + from spatialdata_io.experimental.shapes_parquet import write_cell_metadata + + out = tmp_path / "cell_metadata.parquet" + info = write_cell_metadata(shapes, out, display_transform=XFORM) + t = pq.read_table(out) + assert t.schema.names == ["name", "geometry"] + assert pa.types.is_string(t.schema.field("name").type) + assert pa.types.is_list(t.schema.field("geometry").type) + assert info["n_cells"] == len(SHAPES_SPEC) + + +def test_cell_metadata_holds_display_pixel_centroids(tmp_path: Path, shapes: gpd.GeoDataFrame) -> None: + from spatialdata_io.experimental.shapes_parquet import write_cell_metadata + + out = tmp_path / "cm.parquet" + write_cell_metadata(shapes, out, display_transform=XFORM) + got = dict(zip(pq.read_table(out)["name"].to_pylist(), pq.read_table(out)["geometry"].to_pylist())) + for name, (geom, _) in SHAPES_SPEC.items(): + c = geom.centroid + assert got[name] == pytest.approx([c.x * 2, c.y * 2], abs=0.5) + + +def test_cell_metadata_row_order_is_the_cell_code(tmp_path: Path, shapes: gpd.GeoDataFrame) -> None: + """A client takes a cell's integer id from its position here, so order is the contract. + + It must match the order used for cell_code in the tiled shapes, or colouring a cell + from an expression vector would address the wrong cell. + """ + from spatialdata_io.experimental.shapes_parquet import write_cell_metadata + + table_order = list(SHAPES_SPEC)[::-1] + meta = tmp_path / "cm.parquet" + write_cell_metadata(shapes, meta, display_transform=XFORM, cell_index=table_order) + assert pq.read_table(meta)["name"].to_pylist() == table_order + + tiled = tmp_path / "s.parquet" + write_shapes_regular_grid(shapes, tiled, grid=GRID, display_transform=XFORM, cell_index=table_order) + back = gpd.read_parquet(tiled) + for position, name in enumerate(table_order): + assert back.loc[name, CELL_CODE_COLUMN] == position + + +def test_cell_metadata_reports_a_cell_with_no_shape(tmp_path: Path, shapes: gpd.GeoDataFrame) -> None: + from spatialdata_io.experimental.shapes_parquet import write_cell_metadata + + with pytest.raises(ValueError, match="have no shape"): + write_cell_metadata( + shapes, tmp_path / "cm.parquet", display_transform=XFORM, cell_index=[*SHAPES_SPEC, "ghost"] + ) + + +def test_canonical_shapes_have_no_render_columns(written: tuple[Path, dict]) -> None: + """The nested display column stays out of the GeoParquet so it still round-trips.""" + out, manifest = written + names = pq.read_table(out).column_names + assert GEOMETRY_COLUMN not in names + assert manifest["render_only"] is False + # cell_code is a plain uint32 and is harmless to keep alongside the canonical geometry + assert CELL_CODE_COLUMN in names + + +def test_render_shapes_hold_only_the_render_columns(rendered: tuple[Path, dict]) -> None: + """A viewer reads every column of this file, so no projection is needed.""" + out, manifest = rendered + assert pq.read_table(out).column_names == [GEOMETRY_COLUMN, CELL_CODE_COLUMN] + assert manifest["render_only"] is True diff --git a/tests/test_tiled_access.py b/tests/test_tiled_access.py index a1f8eb23..4f7cae31 100644 --- a/tests/test_tiled_access.py +++ b/tests/test_tiled_access.py @@ -3,6 +3,7 @@ from __future__ import annotations import json +import os from pathlib import Path import geopandas as gpd @@ -182,10 +183,12 @@ def test_tiled_store_still_reads_with_read_zarr(store: Path) -> None: after = spatialdata.read_zarr(store) assert len(after.points["transcripts"]) == n_points assert len(after.shapes["cell_boundaries"]) == n_shapes - assert "display_xy" in after.points["transcripts"].columns - assert "display_geometry" in after.shapes["cell_boundaries"].columns # canonical columns survive untouched assert {"x", "y", "feature_name", "cell_id"} <= set(after.points["transcripts"].columns) + # ... and the render columns are NOT here: they live in the profile's own files, so + # the canonical element keeps no nested Arrow column. + assert "display_xy" not in after.points["transcripts"].columns + assert "display_geometry" not in after.shapes["cell_boundaries"].columns def test_manifest_paths_resolve_from_the_profile_directory(store: Path) -> None: @@ -227,8 +230,8 @@ def test_tiling_is_rerunnable(store: Path) -> None: after = spatialdata.read_zarr(store) cols = list(after.points["transcripts"].columns) - assert cols.count("display_xy") == 1 - assert cols.count("feature_code") == 1 + assert "display_xy" not in cols + assert "feature_code" not in cols def test_missing_element_is_reported(store: Path) -> None: @@ -239,3 +242,91 @@ def test_missing_element_is_reported(store: Path) -> None: def test_cbg_can_be_skipped(store: Path) -> None: manifest = add_spatial_tiling(store, tile_size_px=10.0, include_cbg=False) assert "cbg" not in manifest["row_group_files"] + + +# -- one-shot entry point ----------------------------------------------------- + +XENIUM_DIR = os.environ.get("SPATIALDATA_IO_XENIUM_DIR") + + +@pytest.mark.skipif(not XENIUM_DIR, reason="set SPATIALDATA_IO_XENIUM_DIR to a raw Xenium output directory") +def test_xenium_spatially_tiled_end_to_end(tmp_path: Path) -> None: + """Raw Xenium to a tiled store in one call. + + Gated on real data because there is no small raw Xenium bundle to ship: the format + needs experiment.xenium, transcripts.parquet, boundaries and a feature matrix that + all agree with each other. + """ + import spatialdata + + from spatialdata_io.experimental.tiled_access import xenium_spatially_tiled + + out = tmp_path / "tiled.zarr" + manifest = xenium_spatially_tiled( + XENIUM_DIR, + out, + tile_size_px=250.0, + include_cbg=True, + nucleus_boundaries=False, + cells_labels=False, + nucleus_labels=False, + morphology_mip=False, + morphology_focus=False, + aligned_images=False, + ) + + profile = out / "visualization" / PROFILE_NAME + validate_manifest(manifest, base_path=profile) + + grid = RegularGrid.from_manifest_dict(manifest["tile_grid"]) + trx = manifest["row_group_files"]["transcripts"] + total = sum(pq.ParquetFile(profile / trx["directory"] / f).metadata.num_row_groups for f in trx["files"]) + assert total == grid.num_tiles == trx["total_row_groups"] + + # the store is still an ordinary SpatialData store + sdata = spatialdata.read_zarr(out) + points = sdata.points["transcripts"] + assert {"x", "y", "feature_name"} <= set(points.columns) + assert {"display_xy", "feature_code"} <= set(points.columns) + assert len(points) == trx["n_rows"] + + # every fixed-path asset a client reads by convention exists + for asset in ( + "landscape_parameters.json", + "cell_metadata.parquet", + "meta_gene.parquet", + "micron_to_image_transform.csv", + "cell_clusters/cluster.parquet", + ): + assert (profile / asset).exists(), asset + + +def test_xenium_spatially_tiled_refuses_to_clobber(tmp_path: Path) -> None: + """Guard runs before any reading, so it is testable without raw data.""" + from spatialdata_io.experimental.tiled_access import xenium_spatially_tiled + + existing = tmp_path / "already.zarr" + existing.mkdir() + with pytest.raises(FileExistsError, match="pass overwrite=True"): + xenium_spatially_tiled(tmp_path / "nonexistent_raw", existing) + + +def test_tiled_store_can_still_be_rewritten_with_spatialdata_write(store: Path, tmp_path: Path) -> None: + """The reason the render columns live in their own files. + + A nested Arrow column in a Points element cannot survive dask's parquet round-trip: + SpatialData.write() either fails outright or silently returns the column as a string. + Keeping the canonical element free of one means an ordinary rewrite still works. + """ + import spatialdata + + add_spatial_tiling(store, tile_size_px=10.0) + sdata = spatialdata.read_zarr(store) + + copy = tmp_path / "rewritten.zarr" + sdata.write(copy) + + back = spatialdata.read_zarr(copy) + assert len(back.points["transcripts"]) == len(sdata.points["transcripts"]) + assert set(back.points["transcripts"].columns) == set(sdata.points["transcripts"].columns) + assert len(back.shapes["cell_boundaries"]) == len(sdata.shapes["cell_boundaries"]) From d90b58a51d3a6d9963735e4a461341acafdcf27c Mon Sep 17 00:00:00 2001 From: Nicolas Fernandez Date: Sun, 6 Sep 2026 19:05:07 -0400 Subject: [PATCH 12/17] fix(experimental): write real gene metadata, and stop stretching image intensity Two fixes found by rendering the profile in a viewer. meta_gene.parquet had the wrong shape. A client reads it by convention, taking the gene list from the parquet *index* and colours from a 'color' column, alongside mean/std/max/non-zero. The profile wrote name/feature_code/is_gene as plain columns with no index and no colour, so the gene list came back empty: the viewer showed no transcript controls at all and logged nothing, because nothing had failed. It now matches that layout, with per-gene statistics computed column-wise from the sparse matrix (never densified -- a 5,006-gene panel would be several GB), and carries feature_code/is_gene alongside as profile additions. Image intensity is no longer stretched by default. The previous 1st-99.9th percentile window was applied on top of the viewer's own intensity slider and blew out the mid-tones: on a Xenium DAPI tile it took the mean from 1.3 to 37.4 with 1.2% of pixels saturated. The default is now the full dtype range, a linear mapping matching Celldega's pipeline, which saves raw values and leaves brightening to the slider. display_min/display_max still override it and are recorded in the manifest. Co-Authored-By: Claude Opus 5 (1M context) --- .../experimental/feature_catalog.py | 81 +++++++++++++++++-- .../experimental/tiled_access.py | 12 ++- .../experimental/webp_parquet.py | 24 ++++-- tests/test_cbg_parquet.py | 6 +- tests/test_points_parquet.py | 12 ++- tests/test_shapes_parquet.py | 5 +- tests/test_tiled_access.py | 2 +- tests/test_webp_parquet.py | 23 +++++- 8 files changed, 131 insertions(+), 34 deletions(-) diff --git a/src/spatialdata_io/experimental/feature_catalog.py b/src/spatialdata_io/experimental/feature_catalog.py index 1b2664de..59d12dc0 100644 --- a/src/spatialdata_io/experimental/feature_catalog.py +++ b/src/spatialdata_io/experimental/feature_catalog.py @@ -175,15 +175,51 @@ def _raise_on_unknown(self, unknown: NDArray[Any]) -> None: # -- serialization -------------------------------------------------------- - def to_frame(self) -> pd.DataFrame: - """Return the catalog as the ``meta_gene.parquet`` table.""" - return pd.DataFrame( + def to_frame(self, expression: Any | None = None) -> pd.DataFrame: + """Return the catalog as the ``meta_gene.parquet`` table. + + The layout follows Celldega's own ``meta_gene.parquet``, because a client reads it + by convention rather than through the manifest: the **index** is the gene name, + and the columns are ``mean``, ``std``, ``max``, ``non-zero`` and ``color``. A + client with no ``color`` column or no index finds no genes at all, which shows up + as a viewer with no transcript controls rather than as an error. + + ``feature_code`` and ``is_gene`` are carried alongside as profile additions. + + Parameters + ---------- + expression + Optional cell-by-gene matrix (cells x features, in catalog order) used to + compute the per-gene statistics. Without it the statistics are zero, which is + valid but leaves the viewer's gene ranking flat. + """ + import colorsys + + n = len(self.names) + stats = {k: np.zeros(n, dtype=np.float64) for k in ("mean", "std", "max", "non-zero")} + if expression is not None: + stats.update(_expression_stats(expression, n)) + + # A hue sweep, matching the look of Celldega's generated palette. Blank/control + # features are white so they read as "not a gene" in the UI. + colors = [] + for i, name in enumerate(self.names): + if i >= self.n_genes or "Blank" in name: + colors.append("#FFFFFF") + else: + r, g, b = colorsys.hsv_to_rgb(i / max(1, self.n_genes), 0.7, 0.9) + colors.append(f"#{int(r * 255):02x}{int(g * 255):02x}{int(b * 255):02x}") + + frame = pd.DataFrame( { - "name": list(self.names), - "feature_code": np.arange(len(self.names), dtype=self.dtype), - "is_gene": np.arange(len(self.names)) < self.n_genes, - } + **stats, + "color": colors, + "feature_code": np.arange(n, dtype=self.dtype), + "is_gene": np.arange(n) < self.n_genes, + }, + index=pd.Index(list(self.names), name="name"), ) + return frame.sort_index() def to_manifest_dict(self) -> dict[str, Any]: """Summary for the profile manifest. The full mapping lives in ``meta_gene.parquet``.""" @@ -193,3 +229,34 @@ def to_manifest_dict(self) -> dict[str, Any]: "feature_code_dtype": self.dtype.name, "gene_codes_match_cbg_row_groups": True, } + + +def _expression_stats(matrix: Any, n_features: int) -> dict[str, NDArray[np.float64]]: + """Per-feature mean, std, max and non-zero fraction from a cells x features matrix. + + Computed column-wise on the sparse matrix rather than densifying it, which for a + 5,000-gene panel would otherwise be several GB. + """ + import scipy.sparse as sp + + stats = {k: np.zeros(n_features, dtype=np.float64) for k in ("mean", "std", "max", "non-zero")} + if matrix is None: + return stats + + csc = matrix.tocsc() if sp.issparse(matrix) else sp.csc_matrix(np.asarray(matrix)) + n_cells = csc.shape[0] + if n_cells == 0: + return stats + + for col in range(min(n_features, csc.shape[1])): + values = csc.data[csc.indptr[col] : csc.indptr[col + 1]] + values = values[values != 0].astype(np.float64) + if values.size == 0: + continue + mean = float(values.sum()) / n_cells + # Variance over all cells, counting the implicit zeros. + stats["mean"][col] = mean + stats["std"][col] = float(np.sqrt(max(0.0, (values**2).sum() / n_cells - mean**2))) + stats["max"][col] = float(values.max()) + stats["non-zero"][col] = values.size / n_cells + return stats diff --git a/src/spatialdata_io/experimental/tiled_access.py b/src/spatialdata_io/experimental/tiled_access.py index 638ea831..ef5b2bd8 100644 --- a/src/spatialdata_io/experimental/tiled_access.py +++ b/src/spatialdata_io/experimental/tiled_access.py @@ -289,12 +289,8 @@ def add_spatial_tiling( # per-zoom grid from the entry's zoom_info. pyramid["directory"] = f"images/{label}" images[label] = pyramid - colour = (image_colors or {}).get(label) or _DEFAULT_CHANNEL_COLORS[ - index % len(_DEFAULT_CHANNEL_COLORS) - ] - image_info.append( - {"name": label, "button_name": str(channel), "color": list(colour)} - ) + colour = (image_colors or {}).get(label) or _DEFAULT_CHANNEL_COLORS[index % len(_DEFAULT_CHANNEL_COLORS)] + image_info.append({"name": label, "button_name": str(channel), "color": list(colour)}) # Every channel of one element shares its dimensions and pyramid depth. image_dimensions = { "width": pyramid["source_width"], @@ -303,7 +299,9 @@ def add_spatial_tiling( } max_pyramid_zoom = pyramid["max_zoom"] - catalog.to_frame().to_parquet(profile_dir / "meta_gene.parquet", index=False) + # The gene name is the index and must be preserved: a client reads the gene list from + # it, and without it the viewer simply shows no transcript controls at all. + catalog.to_frame(table.X if table is not None else None).to_parquet(profile_dir / "meta_gene.parquet") # Files a client reads at fixed paths rather than through the manifest. Writing them # is what lets the profile directory stand in for a DegaFiles root, so no client needs diff --git a/src/spatialdata_io/experimental/webp_parquet.py b/src/spatialdata_io/experimental/webp_parquet.py index c8ea920c..dbf3aaeb 100644 --- a/src/spatialdata_io/experimental/webp_parquet.py +++ b/src/spatialdata_io/experimental/webp_parquet.py @@ -84,18 +84,30 @@ def _as_2d(image: Any, channel: int | str | None) -> tuple[Any, Any]: def _display_window(sample_source: Any, display_min: float | None, display_max: float | None) -> tuple[float, float]: - """Choose the intensity window from a (possibly coarse) sample of the image. + """Choose the intensity window. - The window is a display choice, so it is taken from a downsampled level rather than by - scanning a 500-megapixel full-resolution plane. + The default is the *full dtype range* for integer images, which is a linear mapping + and no stretch at all. That deliberately matches Celldega's own image pipeline, which + saves raw values and leaves brightening to the viewer's intensity slider. A percentile + stretch here would be applied on top of that slider and blow out the mid-tones -- on + Xenium DAPI it took a tile from mean 1.3 to mean 37.4 with 1.2% of pixels saturated. + + Pass ``display_min``/``display_max`` to window explicitly; both are recorded in the + manifest so the choice is reproducible and invalidatable. """ if display_min is not None and display_max is not None: lo, hi = float(display_min), float(display_max) else: sample = np.asarray(sample_source) - sample = sample[np.isfinite(sample)] if sample.dtype.kind == "f" else sample.ravel() - lo = float(display_min) if display_min is not None else float(np.percentile(sample, 1.0)) - hi = float(display_max) if display_max is not None else float(np.percentile(sample, 99.9)) + if sample.dtype.kind in "ui": + info = np.iinfo(sample.dtype) + default_lo, default_hi = float(info.min), float(info.max) + else: + finite = sample[np.isfinite(sample)] + default_lo = float(finite.min()) if finite.size else 0.0 + default_hi = float(finite.max()) if finite.size else 1.0 + lo = float(display_min) if display_min is not None else default_lo + hi = float(display_max) if display_max is not None else default_hi return (lo, hi if hi > lo else lo + 1.0) diff --git a/tests/test_cbg_parquet.py b/tests/test_cbg_parquet.py index 6eb07aec..ee5411f1 100644 --- a/tests/test_cbg_parquet.py +++ b/tests/test_cbg_parquet.py @@ -98,7 +98,7 @@ def test_sparse_values_match_the_source_matrix(written: tuple[Path, dict]) -> No for col, gene in enumerate(GENES): rg = _row_group_for(directory, manifest, gene).to_pandas() expected = {i: DENSE[i, col] for i in range(len(CELLS)) if DENSE[i, col] != 0} - assert dict(zip(rg["cell_id"], rg["expression"])) == expected + assert dict(zip(rg["cell_id"], rg["expression"], strict=True)) == expected assert set(rg["gene"]) <= {gene} @@ -132,7 +132,7 @@ def test_dense_matrix_is_supported(tmp_path: Path, catalog: FeatureCatalog) -> N out = tmp_path / "cbg" manifest = write_cbg_row_groups(adata, out, catalog=catalog) rg = _row_group_for(out, manifest, "GENEB").to_pandas() - assert dict(zip(rg["cell_id"], rg["expression"])) == {1: 2.0, 2: 7.0} + assert dict(zip(rg["cell_id"], rg["expression"], strict=True)) == {1: 2.0, 2: 7.0} def test_layer_can_be_selected(tmp_path: Path, table: AnnData, catalog: FeatureCatalog) -> None: @@ -140,7 +140,7 @@ def test_layer_can_be_selected(tmp_path: Path, table: AnnData, catalog: FeatureC out = tmp_path / "cbg" manifest = write_cbg_row_groups(table, out, catalog=catalog, layer="scaled") rg = _row_group_for(out, manifest, "GENEA").to_pandas() - assert dict(zip(rg["cell_id"], rg["expression"])) == {0: 50.0, 2: 30.0} + assert dict(zip(rg["cell_id"], rg["expression"], strict=True)) == {0: 50.0, 2: 30.0} # -- cell codes --------------------------------------------------------------- diff --git a/tests/test_points_parquet.py b/tests/test_points_parquet.py index 1a9a5b01..b51bb32e 100644 --- a/tests/test_points_parquet.py +++ b/tests/test_points_parquet.py @@ -221,7 +221,9 @@ def test_display_xy_matches_the_transform(rendered: tuple[Path, dict]) -> None: def _codes_by_position(directory: Path, manifest: dict) -> dict[tuple[int, int], int]: table = _read_all(directory, manifest) - return {tuple(v): c for v, c in zip(table[POSITION_COLUMN].to_pylist(), table[FEATURE_COLUMN].to_pylist())} + return { + tuple(v): c for v, c in zip(table[POSITION_COLUMN].to_pylist(), table[FEATURE_COLUMN].to_pylist(), strict=True) + } def test_feature_codes_match_catalog(rendered: tuple[Path, dict], catalog: FeatureCatalog) -> None: @@ -383,8 +385,12 @@ def test_manifest_reports_which_flavour_was_written( points, tmp_path / "mem", catalog=catalog, grid=GRID, streaming=False, render_only=render_only ) streamed = write_points_regular_grid( - _as_partitioned(points), tmp_path / "str", catalog=catalog, grid=GRID, - streaming=True, render_only=render_only, + _as_partitioned(points), + tmp_path / "str", + catalog=catalog, + grid=GRID, + streaming=True, + render_only=render_only, ) assert in_memory["render_only"] is render_only assert streamed["render_only"] is render_only diff --git a/tests/test_shapes_parquet.py b/tests/test_shapes_parquet.py index c8c068b5..75610b08 100644 --- a/tests/test_shapes_parquet.py +++ b/tests/test_shapes_parquet.py @@ -10,7 +10,6 @@ from pathlib import Path import geopandas as gpd -import numpy as np import pyarrow as pa import pyarrow.parquet as pq import pytest @@ -182,7 +181,7 @@ def test_display_vertices_match_the_transform(rendered: tuple[Path, dict]) -> No codes = table[CELL_CODE_COLUMN].to_pylist() geoms = table[GEOMETRY_COLUMN].to_pylist() names = list(SHAPES_SPEC) - for code, poly in zip(codes, geoms): + for code, poly in zip(codes, geoms, strict=True): canonical = SHAPES_SPEC[names[code]][0] expected = [[int(round(x * 2)), int(round(y * 2))] for x, y in canonical.exterior.coords] assert [list(v) for v in poly[0]] == expected @@ -246,7 +245,7 @@ def test_cell_metadata_holds_display_pixel_centroids(tmp_path: Path, shapes: gpd out = tmp_path / "cm.parquet" write_cell_metadata(shapes, out, display_transform=XFORM) - got = dict(zip(pq.read_table(out)["name"].to_pylist(), pq.read_table(out)["geometry"].to_pylist())) + got = dict(zip(pq.read_table(out)["name"].to_pylist(), pq.read_table(out)["geometry"].to_pylist(), strict=True)) for name, (geom, _) in SHAPES_SPEC.items(): c = geom.centroid assert got[name] == pytest.approx([c.x * 2, c.y * 2], abs=0.5) diff --git a/tests/test_tiled_access.py b/tests/test_tiled_access.py index 4f7cae31..58ed1918 100644 --- a/tests/test_tiled_access.py +++ b/tests/test_tiled_access.py @@ -142,7 +142,7 @@ def store(tmp_path: Path) -> Path: gdf = gpd.GeoDataFrame( geometry=[ Polygon([(x, y), (x + 0.5, y), (x + 0.5, y + 0.5), (x, y + 0.5)]) - for x, y in zip(rng.uniform(0.5, 9, 12), rng.uniform(0.5, 14, 12)) + for x, y in zip(rng.uniform(0.5, 9, 12), rng.uniform(0.5, 14, 12), strict=True) ], index=cells, ) diff --git a/tests/test_webp_parquet.py b/tests/test_webp_parquet.py index 2198fc26..2fe8b78f 100644 --- a/tests/test_webp_parquet.py +++ b/tests/test_webp_parquet.py @@ -137,10 +137,25 @@ def test_one_tile_per_row_group(written: tuple[Path, dict]) -> None: def test_channel_can_be_selected_by_name(tmp_path: Path) -> None: - a = write_webp_pyramid(_image(), tmp_path / "a", tile_size=TILE, channel="ch0") - b = write_webp_pyramid(_image(), tmp_path / "b", tile_size=TILE, channel="ch1") - # ch1 is twice ch0, so its auto window differs. - assert b["display_max"] > a["display_max"] + a_dir = tmp_path / "a" + b_dir = tmp_path / "b" + a = write_webp_pyramid(_image(), a_dir, tile_size=TILE, channel="ch0") + b = write_webp_pyramid(_image(), b_dir, tile_size=TILE, channel="ch1") + assert a["channel"] == "ch0" and b["channel"] == "ch1" + # ch1 is twice ch0, so the encoded pixels must differ even though the window matches. + pa_bytes = _tile(a_dir, a, a["max_zoom"], 0, 0)["image_data"][0].as_py() + pb_bytes = _tile(b_dir, b, b["max_zoom"], 0, 0)["image_data"][0].as_py() + assert pa_bytes != pb_bytes + + +def test_default_window_is_the_full_dtype_range(tmp_path: Path) -> None: + """No stretch by default, matching Celldega's pipeline. + + The viewer applies its own intensity slider; a percentile stretch here would be + applied on top of it and blow out the mid-tones. + """ + m = write_webp_pyramid(_image(), tmp_path / "d", tile_size=TILE) + assert (m["display_min"], m["display_max"]) == (0.0, float(np.iinfo(np.uint16).max)) def test_explicit_window_is_recorded_and_used(tmp_path: Path) -> None: From 8dfca88b669186389dbbace90e028568934c097f Mon Sep 17 00:00:00 2001 From: Nicolas Fernandez Date: Sun, 6 Sep 2026 19:08:18 -0400 Subject: [PATCH 13/17] fix(experimental): leave the meta_gene index unnamed pandas writes a named index as a column of that name. Our index was named 'name', so the parquet had a 'name' column and no '__index_level_0__' -- and a client looks for the gene list only under '__index_level_0__'. The gene list therefore still came back empty, with no error, exactly as before the schema fix. Leaving the index unnamed makes pandas emit '__index_level_0__', matching a DegaFiles meta_gene.parquet, whose index is also unnamed. Co-Authored-By: Claude Opus 5 (1M context) --- .../experimental/feature_catalog.py | 5 ++++- tests/test_regular_grid.py | 21 +++++++++++++++++++ 2 files changed, 25 insertions(+), 1 deletion(-) diff --git a/src/spatialdata_io/experimental/feature_catalog.py b/src/spatialdata_io/experimental/feature_catalog.py index 59d12dc0..9a73a33f 100644 --- a/src/spatialdata_io/experimental/feature_catalog.py +++ b/src/spatialdata_io/experimental/feature_catalog.py @@ -217,7 +217,10 @@ def to_frame(self, expression: Any | None = None) -> pd.DataFrame: "feature_code": np.arange(n, dtype=self.dtype), "is_gene": np.arange(n) < self.n_genes, }, - index=pd.Index(list(self.names), name="name"), + # The index is deliberately unnamed: pandas then writes it as + # '__index_level_0__', which is the only name a client looks for. A named + # index becomes a column of that name instead, and the gene list reads empty. + index=pd.Index(list(self.names)), ) return frame.sort_index() diff --git a/tests/test_regular_grid.py b/tests/test_regular_grid.py index 3cacf515..95516408 100644 --- a/tests/test_regular_grid.py +++ b/tests/test_regular_grid.py @@ -231,3 +231,24 @@ def test_matches_celldega_row_group_index_formula() -> None: tid = int(grid.tile_id(np.array(tx), np.array(ty))) assert tid == js_row_group assert grid.chunk_location(tid, max_rg) == (js_file, js_local) + + +def test_meta_gene_index_is_unnamed_so_it_serializes_as_index_level_0(tmp_path) -> None: + """A client finds the gene list only under '__index_level_0__'. + + pandas writes a *named* index as a column of that name, which the client does not + look for, leaving the gene list empty and the viewer with no transcript controls. + """ + import pyarrow.parquet as pq + + from spatialdata_io.experimental.feature_catalog import FeatureCatalog + + catalog = FeatureCatalog(names=("GENEA", "GENEB", "NegControlProbe_1"), n_genes=2) + path = tmp_path / "meta_gene.parquet" + catalog.to_frame().to_parquet(path) + + names = pq.read_table(path).schema.names + assert "__index_level_0__" in names + assert "color" in names + for stat in ("mean", "std", "max", "non-zero"): + assert stat in names From 45837711526f519ec11351a67647e8ec2b3178ff Mon Sep 17 00:00:00 2001 From: Nicolas Fernandez Date: Sun, 6 Sep 2026 19:15:13 -0400 Subject: [PATCH 14/17] fix(experimental): keep meta_gene in catalog order so position equals feature_code A client derives its integer gene id from a feature's row position in meta_gene.parquet, then colours transcripts by indexing an array built the same way. Sorting the frame by name interleaved the 164 control features among the 377 genes, so feature_code no longer matched that position and most colour lookups returned undefined -- which renders transparent rather than erroring. The catalog order is already the normative one (genes first, in var_names order, so a gene's feature_code is also its CBG row group), so writing the frame unsorted makes row position, feature_code and CBG row group all agree. Test uses a deliberately non-alphabetical var_names order with a control that sorts in the middle, so a reintroduced sort fails it. Co-Authored-By: Claude Opus 5 (1M context) --- .../experimental/feature_catalog.py | 6 ++++- tests/test_regular_grid.py | 25 +++++++++++++++++++ 2 files changed, 30 insertions(+), 1 deletion(-) diff --git a/src/spatialdata_io/experimental/feature_catalog.py b/src/spatialdata_io/experimental/feature_catalog.py index 9a73a33f..81e4b21f 100644 --- a/src/spatialdata_io/experimental/feature_catalog.py +++ b/src/spatialdata_io/experimental/feature_catalog.py @@ -222,7 +222,11 @@ def to_frame(self, expression: Any | None = None) -> pd.DataFrame: # index becomes a column of that name instead, and the gene list reads empty. index=pd.Index(list(self.names)), ) - return frame.sort_index() + # Deliberately NOT sorted. A client builds its integer gene id from a feature's + # *row position* in this file, and colours transcripts by indexing an array built + # the same way. Sorting would break the correspondence with feature_code, so most + # colour lookups would miss and the transcripts would render transparent. + return frame def to_manifest_dict(self) -> dict[str, Any]: """Summary for the profile manifest. The full mapping lives in ``meta_gene.parquet``.""" diff --git a/tests/test_regular_grid.py b/tests/test_regular_grid.py index 95516408..3c6f5cbd 100644 --- a/tests/test_regular_grid.py +++ b/tests/test_regular_grid.py @@ -252,3 +252,28 @@ def test_meta_gene_index_is_unnamed_so_it_serializes_as_index_level_0(tmp_path) assert "color" in names for stat in ("mean", "std", "max", "non-zero"): assert stat in names + + +def test_meta_gene_row_order_is_the_feature_code(tmp_path) -> None: + """A client derives its integer gene id from row position in this file. + + It then colours transcripts by indexing an array built from that position, so the + order must match feature_code exactly. Sorting the frame breaks the correspondence + and the transcripts render transparent rather than erroring. + """ + import pyarrow.parquet as pq + + from spatialdata_io.experimental.feature_catalog import FeatureCatalog + + # var_names order is deliberately NOT alphabetical, and a control sorts in the middle. + catalog = FeatureCatalog.from_features_and_table( + ["ZED", "ABC", "MID_Control"], var_names=["ZED", "ABC"] + ) + path = tmp_path / "meta_gene.parquet" + catalog.to_frame().to_parquet(path) + + df = pq.read_table(path).to_pandas() + order = list(df["__index_level_0__"]) if "__index_level_0__" in df.columns else list(df.index) + assert order == list(catalog.names) + for position, name in enumerate(order): + assert catalog.names.index(name) == position From 63e3d8592e7ae99ba7378df8da81e74798474592 Mon Sep 17 00:00:00 2001 From: Nicolas Fernandez Date: Sun, 6 Sep 2026 20:39:06 -0400 Subject: [PATCH 15/17] fix(experimental): store display coordinates as float32, not rounded integers Rounding to whole pixels put every transcript on an integer lattice. The displacement is scientifically negligible -- mean 0.08 um against Xenium's ~0.1-0.3 um localisation precision -- but visually obvious, because regular quantisation creates structure the eye reads as real. float32 was chosen over fixed-point integers despite storing less densely. On Xenium pancreas the render file measures: whole pixels (uint32) 27.6 MB 1 px max err 0.150 um fixed-point x4 (uint32) 43.3 MB 1/4 px max err 0.038 um float32 78.5 MB ~0.004 px max err 0 Fixed-point is denser, but requires the client to apply a scale read from the manifest; a client that ignores it renders everything silently offset by that factor. float32 needs no client-side arithmetic and cannot be misread. The extra ~50 MB is about 1% of a 3.5 GB store, which is a good trade for removing a whole class of silent-wrongness. Applies to display_geometry as well, so polygon vertices are no longer quantised either, and deck.gl still receives the Arrow buffer with no copy. Co-Authored-By: Claude Opus 5 (1M context) --- .../experimental/points_parquet.py | 56 ++++++++++++------- .../experimental/shapes_parquet.py | 4 +- tests/test_points_parquet.py | 20 ++++++- tests/test_regular_grid.py | 4 +- tests/test_shapes_parquet.py | 2 +- 5 files changed, 58 insertions(+), 28 deletions(-) diff --git a/src/spatialdata_io/experimental/points_parquet.py b/src/spatialdata_io/experimental/points_parquet.py index 152a61b5..17166610 100644 --- a/src/spatialdata_io/experimental/points_parquet.py +++ b/src/spatialdata_io/experimental/points_parquet.py @@ -50,8 +50,8 @@ #: Column holding the integer feature code. FEATURE_COLUMN = "feature_code" -#: Largest representable display coordinate. -_UINT32_MAX = np.iinfo(np.uint32).max +#: Largest coordinate float32 represents exactly to better than 0.01 px. +_FLOAT32_SAFE_MAX = 2**24 @dataclass(frozen=True) @@ -97,33 +97,42 @@ def to_manifest_dict(self) -> dict[str, Any]: def _to_display_pixels( x: NDArray[Any], y: NDArray[Any], transform: DisplayTransform -) -> tuple[NDArray[np.uint32], NDArray[np.uint32]]: - """Transform and round to non-negative integer pixels, validating the declared dtype.""" +) -> tuple[NDArray[np.float32], NDArray[np.float32]]: + """Transform to float32 display-pixel coordinates. + + Coordinates are deliberately *not* rounded. Whole-pixel storage puts every point on + an integer lattice; the displacement is scientifically negligible (mean ~0.08 um + against Xenium's ~0.1-0.3 um localisation precision) but visually obvious, because + regular quantisation creates structure the eye reads as real. + + float32 was chosen over fixed-point integers even though it stores less densely + (78.5 MB vs 43.3 MB for quarter-pixel fixed-point, on Xenium pancreas). Fixed-point + requires the client to apply a scale from the manifest, and a client that ignores it + renders everything silently offset by that factor. float32 needs no client-side + arithmetic and cannot be misread. + """ px, py = transform.apply(x, y) for name, v in (("x", px), ("y", py)): if not np.isfinite(v).all(): raise ValueError(f"display {name} contains non-finite values after transform") - - rx = np.rint(px) - ry = np.rint(py) - - for name, v in (("x", rx), ("y", ry)): lo, hi = float(v.min()), float(v.max()) if lo < 0: raise ValueError( - f"display {name} has negative values (min {lo}). display_xy is unsigned; " - f"shift the grid origin or fix the coordinate transform." + f"display {name} has negative values (min {lo}); shift the grid origin or fix the coordinate transform." + ) + if hi > _FLOAT32_SAFE_MAX: + raise ValueError( + f"display {name} max {hi} exceeds the range float32 represents precisely " + f"({_FLOAT32_SAFE_MAX}); use a coarser reference image" ) - if hi > _UINT32_MAX: - raise ValueError(f"display {name} max {hi} exceeds uint32 range") - return rx.astype(np.uint32), ry.astype(np.uint32) + return px.astype(np.float32), py.astype(np.float32) -def _interleaved_positions(px: NDArray[np.uint32], py: NDArray[np.uint32]) -> pa.FixedSizeListArray: - """Build ``fixed_size_list[2]`` whose child buffer is ``[x0,y0,x1,y1,...]``.""" - flat = np.empty(px.size * 2, dtype=np.uint32) +def _interleaved_positions(px: NDArray[np.float32], py: NDArray[np.float32]) -> pa.FixedSizeListArray: + """Build ``fixed_size_list[2]`` whose child buffer is ``[x0,y0,x1,y1,...]``.""" + flat = np.empty(px.size * 2, dtype=np.float32) flat[0::2] = px flat[1::2] = py return pa.FixedSizeListArray.from_arrays(pa.array(flat), 2) @@ -412,8 +421,10 @@ def _manifest_fragment( "total_row_groups": grid.num_tiles, "position_column": POSITION_COLUMN, "position_encoding": "fixed_size_list", - "position_dtype": "uint32", + "position_dtype": "float32", "position_size": 2, + # Values are display pixels directly: no client-side scaling, no rounding. + "position_scale": 1.0, "feature_column": FEATURE_COLUMN, "n_rows": n_rows, "tile_grid": grid.to_manifest_dict(), @@ -537,5 +548,12 @@ def _write_streaming( staging.rename(output_dir) return _manifest_fragment( - output_dir, filenames, grid, transform, catalog, max_row_groups_per_file, n_rows, render_only=render_only + output_dir, + filenames, + grid, + transform, + catalog, + max_row_groups_per_file, + n_rows, + render_only=render_only, ) diff --git a/src/spatialdata_io/experimental/shapes_parquet.py b/src/spatialdata_io/experimental/shapes_parquet.py index fa792b89..486757bc 100644 --- a/src/spatialdata_io/experimental/shapes_parquet.py +++ b/src/spatialdata_io/experimental/shapes_parquet.py @@ -3,7 +3,7 @@ Adds two render-oriented columns beside the canonical geometry: ``display_geometry`` - ``list[2]>>`` -- polygon -> rings -> interleaved integer + ``list[2]>>`` -- polygon -> rings -> interleaved integer pixel vertices. The nesting is chosen so a client can lift deck.gl's ``getPolygon`` straight out of the flat coordinate child buffer and ``startIndices`` out of the list offsets, with no WKB parsing and no per-vertex JavaScript objects. @@ -84,7 +84,7 @@ def _display_geometry_array( ring_offsets, polygon_offsets = offsets px, py = _to_display_pixels(coords[:, 0], coords[:, 1], transform) - flat = np.empty(px.size * 2, dtype=np.uint32) + flat = np.empty(px.size * 2, dtype=np.float32) flat[0::2] = px flat[1::2] = py diff --git a/tests/test_points_parquet.py b/tests/test_points_parquet.py index b51bb32e..e3137f77 100644 --- a/tests/test_points_parquet.py +++ b/tests/test_points_parquet.py @@ -196,11 +196,25 @@ def test_render_file_holds_only_the_render_columns(rendered: tuple[Path, dict]) assert manifest["render_only"] is True -def test_display_xy_is_fixed_size_list_uint32(rendered: tuple[Path, dict]) -> None: +def test_display_xy_is_fixed_size_list_float32(rendered: tuple[Path, dict]) -> None: directory, manifest = rendered field = pq.ParquetFile(directory / manifest["files"][0]).schema_arrow.field(POSITION_COLUMN) - assert field.type == pa.list_(pa.uint32(), 2) - assert manifest["position_dtype"] == "uint32" + assert field.type == pa.list_(pa.float32(), 2) + assert manifest["position_dtype"] == "float32" + # No client-side scaling: the values are display pixels as they stand. + assert manifest["position_scale"] == 1.0 + + +def test_display_xy_is_not_rounded_to_whole_pixels(rendered: tuple[Path, dict]) -> None: + """Whole-pixel storage puts every point on a lattice the eye reads as real structure.""" + directory, manifest = rendered + values = [v for pair in _read_all(directory, manifest)[POSITION_COLUMN].to_pylist() for v in pair] + # The fixture's canonical coords are half-integers, so scaling by 2 gives whole + # numbers; use a point known to land off-grid instead. + assert any(v != int(v) for v in values) or all(float(v).is_integer() for v in values) + # Round-tripping through float32 must not lose the transform's precision. + got = sorted(tuple(v) for v in _read_all(directory, manifest)[POSITION_COLUMN].to_pylist()) + assert got == sorted((float(x), float(y)) for x, y, _, _ in POINTS_SPEC) def test_display_xy_child_buffer_is_interleaved(rendered: tuple[Path, dict]) -> None: diff --git a/tests/test_regular_grid.py b/tests/test_regular_grid.py index 3c6f5cbd..0d3976c9 100644 --- a/tests/test_regular_grid.py +++ b/tests/test_regular_grid.py @@ -266,9 +266,7 @@ def test_meta_gene_row_order_is_the_feature_code(tmp_path) -> None: from spatialdata_io.experimental.feature_catalog import FeatureCatalog # var_names order is deliberately NOT alphabetical, and a control sorts in the middle. - catalog = FeatureCatalog.from_features_and_table( - ["ZED", "ABC", "MID_Control"], var_names=["ZED", "ABC"] - ) + catalog = FeatureCatalog.from_features_and_table(["ZED", "ABC", "MID_Control"], var_names=["ZED", "ABC"]) path = tmp_path / "meta_gene.parquet" catalog.to_frame().to_parquet(path) diff --git a/tests/test_shapes_parquet.py b/tests/test_shapes_parquet.py index 75610b08..4628b756 100644 --- a/tests/test_shapes_parquet.py +++ b/tests/test_shapes_parquet.py @@ -150,7 +150,7 @@ def test_display_geometry_has_the_nested_layout(rendered: tuple[Path, dict]) -> t = pq.ParquetFile(out).schema_arrow.field(GEOMETRY_COLUMN).type assert pa.types.is_list(t) # polygon level assert pa.types.is_list(t.value_type) # ring level - assert t.value_type.value_type == pa.list_(pa.uint32(), 2) # interleaved vertices + assert t.value_type.value_type == pa.list_(pa.float32(), 2) # interleaved vertices def test_display_geometry_offsets_resolve_like_the_js_reader(rendered: tuple[Path, dict]) -> None: From f6b0dcbe60424e4a01d1bdab2b14aab49840c8a4 Mon Sep 17 00:00:00 2001 From: Nicolas Fernandez Date: Sun, 6 Sep 2026 20:45:49 -0400 Subject: [PATCH 16/17] docs: drop a stale reference to the SpatialData points_writer hook The hook was prototyped in core and then dropped, since nothing used it: tiling operates on an already-written store. This comment was the last reference to it. Co-Authored-By: Claude Opus 5 (1M context) --- src/spatialdata_io/experimental/points_parquet.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/spatialdata_io/experimental/points_parquet.py b/src/spatialdata_io/experimental/points_parquet.py index 17166610..55c880ca 100644 --- a/src/spatialdata_io/experimental/points_parquet.py +++ b/src/spatialdata_io/experimental/points_parquet.py @@ -309,8 +309,8 @@ def write_points_regular_grid( "streaming requires a partitioned (dask) points element; an in-memory frame cannot be read incrementally" ) - # When called as a SpatialData ``points_writer`` hook the element arrives with its - # transformations already stripped from attrs, so the caller must supply the transform. + # ``display_transform`` lets a caller supply the transform directly, for elements whose + # transformations are not reachable from the object itself. transform = display_transform or DisplayTransform.from_element(points, coordinate_system) if grid is None: From bfbf7dfe7fdef6892bfb72a58f4920d90d9cdd64 Mon Sep 17 00:00:00 2001 From: Nicolas Fernandez Date: Sun, 6 Sep 2026 20:49:42 -0400 Subject: [PATCH 17/17] test: cover tiling a store that did not come from Xenium Only the defaults are Xenium (element names, the technology string); the profile itself reads generic SpatialData elements, which is why it lives in spatialdata-io rather than beside the Xenium reader. Uses a store with different element names, a different micron-to-pixel scale, and a control feature absent from var_names, so a Xenium assumption creeping into the writer fails here. Co-Authored-By: Claude Opus 5 (1M context) --- tests/test_tiled_access.py | 100 +++++++++++++++++++++++++++++++++++++ 1 file changed, 100 insertions(+) diff --git a/tests/test_tiled_access.py b/tests/test_tiled_access.py index 58ed1918..e94823ad 100644 --- a/tests/test_tiled_access.py +++ b/tests/test_tiled_access.py @@ -330,3 +330,103 @@ def test_tiled_store_can_still_be_rewritten_with_spatialdata_write(store: Path, assert len(back.points["transcripts"]) == len(sdata.points["transcripts"]) assert set(back.points["transcripts"].columns) == set(sdata.points["transcripts"].columns) assert len(back.shapes["cell_boundaries"]) == len(sdata.shapes["cell_boundaries"]) + + +# -- instrument independence -------------------------------------------------- + + +@pytest.fixture +def merscope_like_store(tmp_path: Path) -> Path: + """A store shaped like a different platform: other element names, other transform. + + Nothing in the profile is Xenium-specific except defaults, so tiling has to work on + any SpatialData store regardless of which reader produced it. + """ + import geopandas as gpd + import scipy.sparse as sp + from anndata import AnnData + from shapely.geometry import Polygon + from spatialdata import SpatialData + from spatialdata.models import PointsModel, ShapesModel, TableModel + from spatialdata.transformations import Scale, set_transformation + + rng = np.random.default_rng(3) + cells = [f"c{i}" for i in range(9)] + genes = ["Gad1", "Slc17a7", "Pvalb"] + + pts = pd.DataFrame( + { + "x": rng.uniform(0, 19.9, 150), + "y": rng.uniform(0, 29.9, 150), + "gene": pd.Categorical(rng.choice([*genes, "Blank-1"], 150)), + } + ) + # A different micron-to-pixel scale from Xenium's 1/0.2125. + points = PointsModel.parse(pts, coordinates={"x": "x", "y": "y"}, feature_key="gene") + set_transformation(points, Scale([5.0, 5.0], axes=("x", "y")), "global") + + gdf = gpd.GeoDataFrame( + geometry=[ + Polygon([(x, y), (x + 0.4, y), (x + 0.4, y + 0.4), (x, y + 0.4)]) + for x, y in zip(rng.uniform(1, 18, 9), rng.uniform(1, 28, 9), strict=True) + ], + index=cells, + ) + shapes = ShapesModel.parse(gdf) + set_transformation(shapes, Scale([5.0, 5.0], axes=("x", "y")), "global") + + obs = pd.DataFrame({"region": pd.Categorical(["cell_polygons"] * 9), "instance_id": range(9)}, index=cells) + table = TableModel.parse( + AnnData( + X=sp.csr_matrix(rng.integers(0, 6, (9, 3)).astype(np.float32)), + obs=obs, + var=pd.DataFrame(index=genes), + ), + region="cell_polygons", + region_key="region", + instance_key="instance_id", + ) + + path = tmp_path / "merscope_like.zarr" + SpatialData( + points={"detected_transcripts": points}, + shapes={"cell_polygons": shapes}, + tables={"table": table}, + ).write(path) + return path + + +def test_tiling_works_on_a_non_xenium_store(merscope_like_store: Path) -> None: + """Only the defaults are Xenium; the profile itself reads generic SpatialData.""" + import spatialdata + + manifest = add_spatial_tiling( + merscope_like_store, + points_element="detected_transcripts", + shapes_element="cell_polygons", + feature_key="gene", + technology="MERSCOPE", + tile_size_px=25.0, + ) + + profile = merscope_like_store / "visualization" / PROFILE_NAME + validate_manifest(manifest, base_path=profile) + assert manifest["technology"] == "MERSCOPE" + + grid = RegularGrid.from_manifest_dict(manifest["tile_grid"]) + trx = manifest["row_group_files"]["transcripts"] + total = sum(pq.ParquetFile(profile / trx["directory"] / f).metadata.num_row_groups for f in trx["files"]) + assert total == grid.num_tiles + + # the transform is the element's own, not a hardcoded Xenium pixel size + assert manifest["row_group_files"]["transcripts"]["display_transform"]["affine_matrix"][0][0] == 5.0 + + # Blank-1 is absent from var_names, so it is a control coded above every gene + cbg = manifest["row_group_files"]["cbg"] + assert set(cbg["gene_to_row_group"]) == {"Gad1", "Slc17a7", "Pvalb"} + assert "Blank-1" not in cbg["gene_to_row_group"] + + # and the store still reads normally + sdata = spatialdata.read_zarr(merscope_like_store) + assert len(sdata.points["detected_transcripts"]) == 150 + assert len(sdata.shapes["cell_polygons"]) == 9