diff --git a/docs/api.md b/docs/api.md index 95ff1650..2071127e 100644 --- a/docs/api.md +++ b/docs/api.md @@ -56,6 +56,23 @@ I/O for the `spatialdata` project. experimental.to_legacy_anndata ``` +### Experimental spatial tiling + +These opt-in functions currently target Xenium-compatible stores. The canonical profile +rewrites Points and Shapes Parquet into spatial row groups, stores Shapes as GeoArrow, +adds table metadata and a gene-major `X_csc` layer, and publishes discovery metadata in +the root Zarr attributes. It creates no separate visualization directory. + +```{eval-rst} +.. currentmodule:: spatialdata_io + +.. autosummary:: + :toctree: generated + + experimental.add_spatial_tiling + experimental.xenium_spatially_tiled +``` + ### Utility functions ```{eval-rst} diff --git a/src/spatialdata_io/experimental/__init__.py b/src/spatialdata_io/experimental/__init__.py index 36dd4d8b..e0643a34 100644 --- a/src/spatialdata_io/experimental/__init__.py +++ b/src/spatialdata_io/experimental/__init__.py @@ -2,6 +2,10 @@ from_legacy_anndata, to_legacy_anndata, ) +from spatialdata_io.experimental.tiled_access import ( + add_spatial_tiling, + xenium_spatially_tiled, +) from spatialdata_io.readers.iss import iss _readers_technologies = [ @@ -14,5 +18,9 @@ "from_legacy_anndata", "to_legacy_anndata", ] +_spatial_tiling = [ + "add_spatial_tiling", + "xenium_spatially_tiled", +] -__all__ = _readers_technologies + _readers_file_types + _converters +__all__ = _readers_technologies + _readers_file_types + _converters + _spatial_tiling diff --git a/src/spatialdata_io/experimental/display_colors.py b/src/spatialdata_io/experimental/display_colors.py new file mode 100644 index 00000000..e6022808 --- /dev/null +++ b/src/spatialdata_io/experimental/display_colors.py @@ -0,0 +1,81 @@ +"""Assign display colours to genes and categorical cell annotations. + +Gene colours use ``uns["gene_colors"]`` in ``var_names`` order. This borrows the +shape of Scanpy's categorical palettes; the gene-specific key and alignment rule +are profile conventions, not an AnnData gene-colour standard. Unstructured lists +are not automatically realigned on variable subset/reorder. Callers must maintain +that association; existing lists are preserved without validation by default. + +Cluster colours use ``uns["_colors"]`` in stored category order. +The palette follows the same golden-ratio hue scheme as Celldega's fallback. +""" + +from __future__ import annotations + +import colorsys +from typing import Any + +__all__ = ["add_gene_colors", "add_cluster_colors", "palette", "GENE_COLORS_KEY"] + +#: ``uns`` key holding one hex colour per gene, in ``var_names`` order. Follows AnnData's +#: palette naming pattern; the gene-to-position association must be maintained explicitly. +GENE_COLORS_KEY = "gene_colors" + +#: Successive hues are separated by the golden ratio, which keeps neighbouring entries +#: visually distinct instead of walking through a smooth ramp where adjacent genes look +#: identical. +_GOLDEN_RATIO_CONJUGATE = 0.618033988749895 + +_SATURATION = 0.65 +_LIGHTNESS = 0.55 + + +def palette(n: int) -> list[str]: + """``n`` visually distinct hex colours, deterministic in ``n`` and position.""" + colors = [] + for i in range(n): + hue = (i * _GOLDEN_RATIO_CONJUGATE) % 1.0 + r, g, b = colorsys.hls_to_rgb(hue, _LIGHTNESS, _SATURATION) + colors.append(f"#{round(r * 255):02x}{round(g * 255):02x}{round(b * 255):02x}") + return colors + + +def add_gene_colors(table: Any, key: str = GENE_COLORS_KEY, overwrite: bool = False) -> str | None: + """Add one hex colour per gene to ``table.uns``, in place. + + Colours are ordered by position in ``var``, which is the order a client indexes with + ``feature_code``. Controls are not in ``var`` and get the client's fallback colour. + + Returns + ------- + The ``uns`` key, or ``None`` when one already exists and ``overwrite`` is False. + """ + if key in table.uns and not overwrite: + return None + table.uns[key] = palette(table.n_vars) + return key + + +def add_cluster_colors(table: Any, column: str, overwrite: bool = False) -> str | None: + """Add ``uns["_colors"]`` for a categorical ``obs`` column, in place. + + Follows the scanpy convention: one colour per category, in category order. + + Returns + ------- + The ``uns`` key, or ``None`` when the column is missing or not categorical. + """ + if column not in table.obs: + return None + + values = table.obs[column] + categories = getattr(getattr(values, "cat", None), "categories", None) + if categories is None: + return None + + key = f"{column}_colors" + if key in table.uns and not overwrite: + return None + + table.uns[key] = palette(len(categories)) + return key diff --git a/src/spatialdata_io/experimental/expression_index.py b/src/spatialdata_io/experimental/expression_index.py new file mode 100644 index 00000000..5f211723 --- /dev/null +++ b/src/spatialdata_io/experimental/expression_index.py @@ -0,0 +1,198 @@ +"""Make gene-major expression reads cheap, without inventing a format. + +``X`` is stored CSR (cell-major), which is right for the usual analysis access pattern of +"give me this cell's profile". A viewer wants column access: "give me this gene across all +cells". Fetching a column from CSR touches every chunk, so a client either downloads the +whole matrix or does nothing lazily at all -- 4.5 MB for Xenium pancreas, 54.8 MB for Prime +skin, and unbounded beyond that. + +Two additions fix it, both plain AnnData that round-trips and that any tool can use: + +``var`` statistics + ``mean``, ``std``, ``max`` and ``non_zero`` per gene. Without these a client has to read + the whole matrix just to populate a gene list, however the matrix is laid out. + +a CSC layer + A gene-major copy of ``X``. One gene becomes ``indptr[g]:indptr[g+1]`` -- a slice of one + or two chunks, about 6,600 non-zeros for skin -- instead of the entire matrix. + +The cost is a second copy of the non-zeros. Historical rebuilt skin measurements were +about 104 MB for this CSC layer versus 150.3 MB for the previous gene-major Parquet. +The layer's contents round-trip as AnnData, but custom chunk sizes need not survive an +ordinary rewrite. Fixed-length chunks target average gene density, not gene boundaries. +""" + +from __future__ import annotations + +from typing import Any + +import numpy as np +from numpy.typing import NDArray + +__all__ = [ + "add_gene_statistics", + "add_csc_layer", + "write_csc_layer", + "GENE_STAT_COLUMNS", + "CSC_LAYER", +] + +#: ``var`` columns written by :func:`add_gene_statistics`. +GENE_STAT_COLUMNS = ("mean", "std", "max", "non_zero") + +#: Layer holding the gene-major copy of ``X``. +CSC_LAYER = "X_csc" + + +def _column_stats(matrix: Any, n_genes: int) -> dict[str, NDArray[np.float64]]: + """Per-gene mean, std, max and non-zero fraction, over all cells. + + Zeros count towards the denominator: they are measurements, not missing data, so ``std`` + is the population standard deviation across every cell and ``non_zero`` is the fraction + of cells with a non-zero value. This matches what ``celldega.pre`` writes, and the + client computes the same figures when these columns are absent. + """ + import scipy.sparse as sp + + n_cells = matrix.shape[0] + csc = matrix.tocsc() if sp.issparse(matrix) else sp.csc_matrix(matrix) + + total = np.asarray(csc.sum(axis=0)).ravel().astype(np.float64) + squared = np.asarray(csc.multiply(csc).sum(axis=0)).ravel().astype(np.float64) + counts = np.diff(csc.indptr).astype(np.float64) + + maxima = np.zeros(n_genes, dtype=np.float64) + for gene in range(csc.shape[1]): + start, end = csc.indptr[gene], csc.indptr[gene + 1] + if end > start: + maxima[gene] = float(csc.data[start:end].max()) + + mean = total / n_cells + # Var[x] = E[x^2] - E[x]^2, clipped because float error can make it slightly negative. + variance = np.maximum(0.0, squared / n_cells - mean**2) + + return { + "mean": mean, + "std": np.sqrt(variance), + "max": maxima, + "non_zero": counts / n_cells, + } + + +def add_gene_statistics(table: Any) -> list[str]: + """Add per-gene summary statistics to ``table.var``, in place. + + Returns + ------- + The column names written. + """ + if table.X is None: + return [] + + stats = _column_stats(table.X, table.n_vars) + for name in GENE_STAT_COLUMNS: + table.var[name] = stats[name] + return list(GENE_STAT_COLUMNS) + + +def add_csc_layer(table: Any, layer: str = CSC_LAYER) -> str | None: + """Add a gene-major (CSC) copy of ``X`` as a layer, in place. + + Returns + ------- + The layer name, or ``None`` when there is nothing to transpose. + """ + import scipy.sparse as sp + + if table.X is None: + return None + if not sp.issparse(table.X): + # A dense X is already randomly addressable by column; a CSC copy would only + # double the storage for no gain. + return None + + table.layers[layer] = table.X.tocsc() + return layer + + +#: Target non-zeros per chunk when none is given: a few genes' worth, so one gene costs one +#: or two chunks instead of a slice of a huge one. +DEFAULT_GENES_PER_CHUNK = 2 + + +def csc_chunk_size(nnz: int, n_genes: int, genes_per_chunk: int = DEFAULT_GENES_PER_CHUNK) -> int: + """Chunk length that keeps a single gene to about one chunk. + + AnnData's default chunking is sized for whole-matrix reads -- 162,948 non-zeros per + chunk for Xenium pancreas, against ~6,915 for one gene. Reading a gene then costs 24x + what it needs. Sizing chunks by the average gene fixes that; the cost is more, smaller + chunks, which compress slightly worse. + """ + if n_genes <= 0 or nnz <= 0: + return max(1, nnz) + per_gene = max(1, nnz // n_genes) + return max(1024, per_gene * genes_per_chunk) + + +def write_csc_layer( + table_path: Any, + csc: Any, + *, + layer: str = CSC_LAYER, + chunk: int | None = None, +) -> dict[str, Any]: + """Write a gene-major matrix into a written table's ``layers``, chunked for gene reads. + + Written directly rather than through AnnData because the chunking is the whole point, + and AnnData sizes chunks for whole-matrix access. + + Parameters + ---------- + table_path + Path of the already-written table group inside the store. + csc + The gene-major matrix. + layer + Layer name under ``layers/``. + chunk + Non-zeros per chunk. Defaults to :func:`csc_chunk_size`. + + Returns + ------- + A description of what was written, for the manifest. + """ + from pathlib import Path + + import zarr + + n_genes = csc.shape[1] + nnz = int(csc.nnz) + chunk = chunk or csc_chunk_size(nnz, n_genes) + + root = zarr.open_group(str(Path(table_path) / "layers"), mode="a") + if layer in root: + del root[layer] + group = root.create_group(layer) + + # AnnData's own encoding, so the layer round-trips as an ordinary csc_matrix. + group.attrs["encoding-type"] = "csc_matrix" + group.attrs["encoding-version"] = "0.1.0" + group.attrs["shape"] = list(csc.shape) + + for name, values, dtype in ( + ("data", csc.data, np.float32), + ("indices", csc.indices, np.int32), + ("indptr", csc.indptr, np.int32), + ): + arr = np.asarray(values, dtype=dtype) + # indptr is one value per gene and always read whole, so it stays a single chunk. + chunks = (len(arr),) if name == "indptr" else (min(chunk, max(1, len(arr))),) + group.create_array(name, shape=arr.shape, chunks=chunks, dtype=dtype)[:] = arr + + # AnnData and SpatialData read through consolidated metadata, so a group added after + # the table was written is invisible to them until the index is refreshed. Browsers + # are unaffected -- zarrita reads each node directly -- which makes this exactly the + # kind of difference that shows up only on the Python side. + zarr.consolidate_metadata(zarr.open_group(str(Path(table_path)), mode="a").store) + + return {"layer": layer, "chunk": chunk, "nnz": nnz, "genes": n_genes} diff --git a/src/spatialdata_io/experimental/feature_catalog.py b/src/spatialdata_io/experimental/feature_catalog.py new file mode 100644 index 00000000..d90e6c82 --- /dev/null +++ b/src/spatialdata_io/experimental/feature_catalog.py @@ -0,0 +1,184 @@ +"""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. Table features come first in ``var_names`` order, so a gene's ``feature_code`` + is its column index in X and the derived CSC layer. No second string join is needed. +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 (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()) + 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_manifest_dict(self) -> dict[str, Any]: + """Legacy catalog summary. The native profile uses var_names plus top-level extra_features.""" + 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/manifest.py b/src/spatialdata_io/experimental/manifest.py new file mode 100644 index 00000000..320f9b91 --- /dev/null +++ b/src/spatialdata_io/experimental/manifest.py @@ -0,0 +1,227 @@ +"""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 the display column names. +Without it a client would have to infer the layout, which is exactly what this profile +exists to avoid. + +Key names follow Celldega's ``landscape_parameters.json``; the adapt_dega reader also +understands native SpatialData sources and display column names (``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__ = [ + "ROOT_MANIFEST_KEY", + "write_root_manifest", + "PROFILE_NAME", + "PROFILE_VERSION", + "MANIFEST_FILENAME", + "build_manifest", + "validate_manifest", + "write_manifest", +] + +PROFILE_NAME = "grid_files_v1" +PROFILE_VERSION = "0.1.0" +MANIFEST_FILENAME = "landscape_parameters.json" + +#: Root Zarr attribute holding the manifest when no profile directory is written. +ROOT_MANIFEST_KEY = "spatial_tiling" + + +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, + image_dimensions: dict[str, Any] | None = None, + max_pyramid_zoom: int | None = None, + fixed_path_assets: dict[str, Any] | None = None, + spatialdata: dict[str, Any] | None = None, + feature_catalog: 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 ------------------------- + # Legacy keys remain compatible with the DegaFiles manifest. The spatialdata + # block and display column names select the native adapter and display schema. + "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 spatialdata: + # Declares that gene/cell metadata, expression and images are read from the store + # itself rather than from files here. A manifest without this block describes a + # self-contained bundle, which is what DegaFiles are. + manifest["spatialdata"] = spatialdata + if feature_catalog: + # `feature_code` indexes the full catalog: genes in `var` order, then controls. + # A client reading gene names from `var` alone sees only the genes, so any control + # transcript would index past the end of its colour table and lose its colour -- + # silently, since it is a small fraction of the points. Recording the names beyond + # `var` is what lets the client rebuild the full ordering. + manifest["feature_catalog"] = feature_catalog + 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: + manifest["max_pyramid_zoom"] = max_pyramid_zoom + 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 ("profile", "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_root_manifest(store: Any, manifest: dict[str, Any], key: str = ROOT_MANIFEST_KEY) -> None: + """Record the manifest in the store's root Zarr attributes instead of a file. + + This removes the need for a ``visualization/`` directory. The manifest describes the + row-group layout of the canonical Parquets and must be regenerated after an ordinary + SpatialData rewrite, which does not preserve that physical layout. + + The block is namespaced under one key so it cannot collide with SpatialData's own + ``spatialdata_attrs``. + """ + import zarr + + root = zarr.open_group(str(store), mode="a") + root.attrs[key] = manifest + + +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/points_parquet.py b/src/spatialdata_io/experimental/points_parquet.py new file mode 100644 index 00000000..82bfe21b --- /dev/null +++ b/src/spatialdata_io/experimental/points_parquet.py @@ -0,0 +1,618 @@ +"""Rewrite a Points element into regular-grid row groups. + +Canonical mode preserves the source columns and index, dropping display columns +left by earlier experimental versions. ``render_only=True`` instead writes a +separate two-column display file: interleaved ``fixed_size_list[2]`` +``display_xy`` and unsigned integer ``feature_code``. Coordinates are unrounded. + +Rows are stably grouped by tile and files are split to limit footer size. The +intended contract is one physical row group per tile, including empty tiles. +The current write_table calls need an explicit row-group size to enforce this +for tiles exceeding PyArrow's default row-group limit. +""" + +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 float32 pixel positions. +POSITION_COLUMN = "display_xy" +#: Column holding the integer feature code. +FEATURE_COLUMN = "feature_code" + +#: Current coordinate magnitude guard; this does not guarantee 0.01 px precision. +_FLOAT32_SAFE_MAX = 2**24 + +# PyArrow caps requested Parquet row groups at 64 Mi rows. Above this, one call to +# ``write_table`` necessarily produces more than one physical row group, which would +# break the profile's ``tile_id == row_group_index`` addressing contract. +_MAX_ROWS_PER_TILE = 64 * 1024 * 1024 + + +@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]: + """Serialize the transform for the manifest, so a client can reproduce the mapping.""" + 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.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") + lo, hi = float(v.min()), float(v.max()) + if lo < 0: + raise ValueError( + 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" + ) + + return px.astype(np.float32), py.astype(np.float32) + + +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) + + +#: 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, + render_only: bool = False, + column_order: list[str] | None = None, +) -> tuple[pa.Table, NDArray[np.int64]]: + """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 + 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) + + if render_only: + # Only the render columns. A viewer reads every column of this file, which is + # why no column projection is needed. Projection is broken in released + # parquet-wasm 0.7.2; Celldega temporarily uses a patched build. + 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 while released parquet-wasm lacks the + # projection fix carried by Celldega's temporary patched build. + # + # 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) + # 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) + + # Put the columns a viewer reads next to each other. parquet-wasm coalesces a + # projection into one byte range spanning the first to the last requested column, so + # any column physically between them is fetched too: reading x, y and feature_name + # costs 27.7 KiB when z sits between them and 19.7 KiB when it does not. Column order + # is therefore load-bearing for read cost, not cosmetic. + if column_order: + leading = [c for c in column_order if c in table.column_names] + rest = [c for c in table.column_names if c not in leading] + table = table.select(leading + rest) + + 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]) + _write_one_row_group(writer, table.slice(start, end - start) if end > start else schema.empty_table()) + + +def _write_one_row_group(writer: pq.ParquetWriter, table: pa.Table) -> None: + """Write exactly one physical Parquet row group, including for an empty tile.""" + n_rows = table.num_rows + if n_rows > _MAX_ROWS_PER_TILE: + raise ValueError( + f"one tile contains {n_rows:,} rows, above Parquet's supported single-row-group " + f"limit of {_MAX_ROWS_PER_TILE:,}; reduce tile_size_px" + ) + writer.write_table(table, row_group_size=max(1, n_rows)) + + +def _validate_physical_row_groups( + directory: Path, + filenames: list[str], + *, + total_row_groups: int, + max_row_groups_per_file: int, +) -> None: + """Verify the file footers match the formula encoded in the manifest.""" + for file_index, filename in enumerate(filenames): + remaining = total_row_groups - file_index * max_row_groups_per_file + expected = min(max_row_groups_per_file, remaining) + actual = pq.ParquetFile(directory / filename).metadata.num_row_groups + if actual != expected: + raise RuntimeError( + f"{filename} contains {actual} physical row groups; expected {expected}. " + "The tile-to-row-group addressing contract would be invalid." + ) + + +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, + *, + 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, + max_row_groups_per_file: int = DEFAULT_MAX_ROW_GROUPS_PER_FILE, + compression: str = "zstd", + streaming: bool | None = None, + render_only: bool = False, + column_order: list[str] | None = None, + 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") + + 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 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" + ) + + # ``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: + 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, + render_only=render_only, + column_order=column_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__}") + + table, tile_ids = _prepare_table( + df, + transform=transform, + catalog=catalog, + feature_key=feature_key, + grid=grid, + categories=None, + render_only=render_only, + column_order=column_order, + ) + table, sorted_tile_ids = _sorted_by_tile(table, tile_ids) + + 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"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(), + } + ) + + try: + 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) + _validate_physical_row_groups( + staging, + filenames, + total_row_groups=grid.num_tiles, + max_row_groups_per_file=max_row_groups_per_file, + ) + except BaseException: + shutil.rmtree(staging, ignore_errors=True) + raise + + 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, + int(table.num_rows), + render_only=render_only, + ) + + +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, + render_only: bool = False, +) -> dict[str, Any]: + fragment: dict[str, Any] = { + "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": "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(), + "display_transform": transform.to_manifest_dict(), + "feature_catalog": catalog.to_manifest_dict(), + "render_only": render_only, + } + return fragment + + +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, + render_only: bool = False, + column_order: list[str] | None = None, +) -> 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, + render_only=render_only, + column_order=column_order, + ) + 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"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(), + } + ) + + # -- 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) + _validate_physical_row_groups( + staging, + filenames, + total_row_groups=grid.num_tiles, + max_row_groups_per_file=max_row_groups_per_file, + ) + 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, + render_only=render_only, + ) diff --git a/src/spatialdata_io/experimental/regular_grid.py b/src/spatialdata_io/experimental/regular_grid.py new file mode 100644 index 00000000..4f82245b --- /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 ``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, +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/src/spatialdata_io/experimental/shapes_parquet.py b/src/spatialdata_io/experimental/shapes_parquet.py new file mode 100644 index 00000000..7f94b7a4 --- /dev/null +++ b/src/spatialdata_io/experimental/shapes_parquet.py @@ -0,0 +1,276 @@ +"""Rewrite a Shapes element into regular-grid row groups. + +Canonical mode preserves geometry and GeoParquet metadata and adds a positional +``cell_code`` column. ``render_only=True`` writes a separate display file with +``display_geometry`` (``list[2]>>``) and ``cell_code``. +The flat vertex buffer avoids WKB parsing and per-vertex JavaScript objects. + +Display geometry is lossy: only the exterior ring of the largest polygon part +is retained. Its centroid assigns the shape to one tile; polygons are not duplicated +across intersected tiles. This is a polygon path, not support for circle Shapes. +""" + +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, + _validate_physical_row_groups, + _write_one_row_group, +) +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 float32-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 float32 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.float32) + 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, geometry_encoding: str = "WKB") -> 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 + + # "geoarrow" writes list>> tagged `geoarrow.polygon`. Celldega + # reads those coordinate buffers without parsing WKB, then materialises the visible + # paths for deck.gl's PathLayer. This encoding therefore makes canonical geometry + # browser-readable without adding a separate display column. + table = _geopandas_to_arrow(shapes, index=None, geometry_encoding=geometry_encoding) + 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, + geometry_encoding: str = "WKB", + 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. + + 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 + Sequence defining ``cell_code`` order, or a mapping from shape instance ID to + absolute table row position. 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) + + # 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] + canonical = shapes.drop(columns=stale) if stale else shapes + table = None if render_only else _canonical_geoparquet_table(canonical, geometry_encoding) + + if cell_index is None: + codes = np.arange(len(shapes), dtype=np.uint32) + else: + from collections.abc import Mapping + + positions = dict(cell_index) if isinstance(cell_index, Mapping) else {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)) + + 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") + 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"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(), + } + ) + + 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 + _write_one_row_group(writer, table.slice(start, end - start) if end > start else schema.empty_table()) + if writer is not None: + writer.close() + validation_dir = staging if not single_file else staging.parent + validation_names = filenames if not single_file else [staging.name] + _validate_physical_row_groups( + validation_dir, + validation_names, + total_row_groups=grid.num_tiles, + max_row_groups_per_file=max_row_groups_per_file, + ) + 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, + "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), + "geometry_is_lossy": render_only, + "tile_grid": grid.to_manifest_dict(), + } + if render_only: + fragment["geometry_note"] = "exterior ring of the largest polygon part; canonical geometry retained" + if single_file: + fragment["path"] = output_path.name + else: + fragment["directory"] = output_path.name + fragment["files"] = filenames + return fragment diff --git a/src/spatialdata_io/experimental/tiled_access.py b/src/spatialdata_io/experimental/tiled_access.py new file mode 100644 index 00000000..f3744c4f --- /dev/null +++ b/src/spatialdata_io/experimental/tiled_access.py @@ -0,0 +1,472 @@ +"""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 prepares table annotations before the initial write, then tiles the +written vector Parquets and writes the tuned CSC buffers. + +The operation is opt-in and modifies the store in place: canonical Parquets are +reordered, Shapes gain cell_code, and the table receives derived annotations and a +CSC layer. Original coordinates and geometries remain authoritative. read_zarr works +unchanged, but an ordinary SpatialData.write does not preserve the display profile +or tile row-group layout. Regenerate after saving or editing source data. + +The multi-asset operation is not transactional. The existing-store entry point deletes +and rewrites the table when expression indexing is requested; the one-shot Xenium entry +point avoids that rewrite. Current support assumes Xenium-compatible instance identifiers, +centroid layout and display coordinates. +""" + +from __future__ import annotations + +import shutil +from pathlib import Path +from typing import Any + +import numpy as np + +from spatialdata_io.experimental.display_colors import add_cluster_colors, add_gene_colors +from spatialdata_io.experimental.expression_index import ( + add_gene_statistics, + write_csc_layer, +) +from spatialdata_io.experimental.feature_catalog import FeatureCatalog +from spatialdata_io.experimental.manifest import ( + PROFILE_NAME, + build_manifest, + validate_manifest, + write_manifest, + write_root_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 ( + CELL_CODE_COLUMN, + write_shapes_regular_grid, +) + +__all__ = ["add_spatial_tiling", "xenium_spatially_tiled"] + +#: Directory inside the store holding derived (non-canonical) profile assets. +PROFILE_DIR = "visualization" + + +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 _cell_positions_for_shapes(table: Any, shapes_element: str) -> dict[Any, int]: + """Map one Shapes element's instance IDs to absolute rows in the annotating table.""" + from spatialdata.models import get_table_keys + + regions, region_key, instance_key = get_table_keys(table) + declared = [regions] if isinstance(regions, str) else list(regions) + # Xenium's table annotates ``cell_labels`` while the boundary Shapes carry the same + # cell instance IDs. When there is only one annotated region, those IDs provide an + # unambiguous bridge to the requested boundaries. With several unrelated regions we + # require an exact region match rather than guessing. + if shapes_element in declared: + selected_regions = {shapes_element} + elif len(declared) == 1: + selected_regions = {declared[0]} + else: + raise ValueError( + f"table does not directly annotate shapes element {shapes_element!r}, and its " + f"declared regions {declared} do not identify one unambiguous instance namespace" + ) + + positions: dict[Any, int] = {} + for row_position, (region, instance_id) in enumerate( + zip(table.obs[region_key], table.obs[instance_key], strict=True) + ): + if region not in selected_regions: + continue + if instance_id in positions: + raise ValueError( + f"table has duplicate {instance_key!r} value {instance_id!r} for region {shapes_element!r}" + ) + positions[instance_id] = row_position + return positions + + +def _prepare_expression_index(table: Any, cluster_column: str | None) -> tuple[Any | None, dict[str, Any]]: + """Mutate a table with display metadata and retain a CSC matrix for tuned writing.""" + import scipy.sparse as sp + + stats = add_gene_statistics(table) + colors = add_gene_colors(table) + cluster_colors = add_cluster_colors(table, cluster_column) if cluster_column else None + csc = table.X.tocsc() if sp.issparse(table.X) else None + description: dict[str, Any] = { + "var_statistics": stats, + **({"gene_colors": colors} if colors else {}), + **({"cluster_colors": cluster_colors} if cluster_colors else {}), + } + return csc, description + + +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", + index_expression: bool = True, + cluster_column: str | None = None, + profile_layout: str = "v1", + compression: str = "zstd", + _expression_index: dict[str, Any] | None = None, +) -> 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 gene ordering and expression. + ``None`` derives codes from categorical observed features, includes their names + in the manifest and does not advertise table-backed viewer components. + coordinate_system + Target coordinate system, assumed to coincide with reference-image pixels. + The current code does not validate that assumption. + 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. + index_expression + Add per-gene statistics to ``var`` and a gene-major (CSC) copy of ``X`` as a layer, + so a client can read one gene without downloading the whole matrix. Costs a second + copy of the non-zeros. + cluster_column + Categorical ``obs`` column to colour cells by. Its palette is written to + ``uns["_colors"]``, the scanpy convention. + profile_layout + ``"v1"`` writes the display Parquets and a manifest file under + ``visualization/grid_files_v1``. ``"canonical"`` writes neither: the canonical + Parquets carry everything a viewer needs. Geometry is written as ``geoarrow`` so a + browser can decode coordinate buffers without parsing WKB; point columns are ordered + so a projection spans no unwanted column; and the manifest goes into the store's root + attributes. No ``visualization/`` directory is created. + 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)) + + canonical_only = profile_layout == "canonical" + if profile_layout not in ("v1", "canonical"): + raise ValueError(f"profile_layout must be 'v1' or 'canonical', not {profile_layout!r}") + + profile_dir = store / PROFILE_DIR / PROFILE_NAME + if not canonical_only: + profile_dir.mkdir(parents=True, exist_ok=True) + + # Reading x, y and feature_name costs 27.7 KiB with z between them and 19.7 KiB + # without, because parquet-wasm coalesces a projection into one contiguous byte range. + render_first = ["x", "y", feature_key] if canonical_only else None + geometry_encoding = "geoarrow" if canonical_only else "WKB" + + # 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 = None + if not canonical_only: + 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. + canonical_points = 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, + column_order=render_first, + overwrite=True, + ) + + if canonical_only: + # Describe the canonical file instead of a display file. Coordinates are two + # separate columns rather than one interleaved column, which is what a client has + # to know to read them: there is no display_xy to fall back on. + transcripts = { + **{k: v for k, v in canonical_points.items() if not k.startswith("position_")}, + "directory": f"points/{points_element}/points.parquet", + "position_encoding": "separate_columns", + "position_columns": ["x", "y"], + "feature_column": feature_key, + "feature_encoding": "dictionary", + "render_only": False, + } + + 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] + shapes_transform = DisplayTransform.from_element(shapes, coordinate_system) + cell_positions = _cell_positions_for_shapes(table, shapes_element) if table is not None else None + if not canonical_only: + cell_segmentation = write_shapes_regular_grid( + shapes, + profile_dir / "cell_seg", + grid=grid, + display_transform=shapes_transform, + cell_index=cell_positions, + max_row_groups_per_file=max_row_groups_per_file, + compression=compression, + render_only=True, + overwrite=True, + ) + canonical_shapes = write_shapes_regular_grid( + shapes, + store / "shapes" / shapes_element / "shapes.parquet", + grid=grid, + display_transform=shapes_transform, + cell_index=cell_positions, + max_row_groups_per_file=max_row_groups_per_file, + compression=compression, + geometry_encoding=geometry_encoding, + overwrite=True, + ) + + if canonical_only: + cell_segmentation = { + **{ + k: v + for k, v in canonical_shapes.items() + if k not in {"directory", "path", "geometry_column", "geometry_is_lossy", "geometry_note"} + }, + "geometry_column": "geometry", + "geometry_encoding": "geoarrow.polygon", + "cell_id_column": CELL_CODE_COLUMN, + "display_transform": shapes_transform.to_manifest_dict(), + "geometry_is_lossy": False, + "render_only": False, + } + if "files" in canonical_shapes: + cell_segmentation["directory"] = f"shapes/{shapes_element}/shapes.parquet" + else: + cell_segmentation["path"] = f"shapes/{shapes_element}/shapes.parquet" + + # Gene-major access and a gene list both require reading every non-zero from a CSR + # matrix. Precomputing the statistics and storing a CSC copy turns "download the + # whole matrix" into "read two chunks", which is what makes this scale. + expression_index: dict[str, Any] | None = _expression_index + if index_expression and table is not None and table_element: + csc, expression_index = _prepare_expression_index(table, cluster_column) + # SpatialData refuses to overwrite an element inside the store it was read from + # (scverse/spatialdata#520), so the table is deleted and rewritten. The statistics + # pass has already materialised X, so nothing is read back from the deleted path. + sdata.delete_element_from_disk(table_element) + sdata.write_element(table_element) + + # Written after the table, and directly, because the chunking is the whole point: + # AnnData sizes chunks for whole-matrix reads, which costs 24x too much per gene. + layer = write_csc_layer(store / "tables" / table_element, csc) if csc is not None else None + if layer: + expression_index["csc"] = layer + + native_components = ["images"] + if table is not None: + native_components[:0] = ["metadata", "cbg"] + spatialdata_manifest: dict[str, Any] = { + "store_url": "." if canonical_only else "../..", + **({} if canonical_only else {"native": native_components}), + **({"table": table_element} if table_element else {}), + **({"cluster_column": cluster_column} if cluster_column else {}), + **({"expression_index": expression_index} if expression_index else {}), + } + + manifest = build_manifest( + grid=grid, + technology=technology, + transcripts=transcripts, + cell_segmentation=cell_segmentation, + # Gene and cell metadata, expression and images are all read from the store itself, + # so the profile declares where the store is rather than duplicating its contents. + feature_catalog={ + "n_genes": catalog.n_genes, + "extra_features": list(catalog.names[catalog.n_genes :]), + **({"names": list(catalog.names)} if table is None else {}), + }, + spatialdata=spatialdata_manifest, + source={ + "store": store.name, + "technology": technology, + "points_element": points_element, + "shapes_element": shapes_element, + "table_element": table_element, + "coordinate_system": coordinate_system, + "tile_size_px": tile_size_px, + }, + ) + if canonical_only: + # The root attribute is a storage/access description, not a Celldega settings + # file. The Celldega adapter supplies its own defaults after discovering this + # profile. Keep the v1 file manifest unchanged for DegaFiles compatibility. + for key in ( + "technology", + "use_row_groups", + "use_int_index", + "segmentation_approach", + "tile_size", + "image_info", + "image_format", + ): + manifest.pop(key, None) + manifest["row_group_files"].pop("images", None) + # Canonical paths are relative to the store root, which is also where this + # manifest lives. Validate before publishing it so a malformed profile cannot + # become a blank viewport in the browser. + validate_manifest(manifest, base_path=store) + write_root_manifest(store, manifest) + else: + 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, + 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. + + 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. + compression + Parquet compression codec. + overwrite + Replace ``output_path`` if it exists. + tiling + Extra keyword arguments for :func:`add_spatial_tiling`, for example + ``{"shapes_element": "cell_polygons"}``. Kept separate from ``xenium_kwargs`` + because the two functions have distinct option sets. + 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) + + tiling_options = dict(tiling or {}) + index_expression = bool(tiling_options.pop("index_expression", True)) + table_element = tiling_options.get("table_element", "table") + cluster_column = tiling_options.get("cluster_column") + + sdata = xenium(raw_path, **xenium_kwargs) + expression_index: dict[str, Any] | None = None + csc = None + if index_expression and table_element is not None: + table = sdata.tables[table_element] + csc, expression_index = _prepare_expression_index(table, cluster_column) + sdata.write(output_path) + + # The one-shot path writes table annotations with the initial store, so it never + # deletes and rewrites that table. Only the CSC buffers are replaced afterward to + # give them the small chunks required for per-gene browser reads. + if csc is not None and table_element is not None: + layer = write_csc_layer(output_path / "tables" / table_element, csc) + expression_index = expression_index or {} + expression_index["csc"] = layer + + return add_spatial_tiling( + output_path, + tile_size_px=tile_size_px, + max_row_groups_per_file=max_row_groups_per_file, + compression=compression, + index_expression=False, + _expression_index=expression_index, + **tiling_options, + ) diff --git a/tests/test_display_colors.py b/tests/test_display_colors.py new file mode 100644 index 00000000..6edb54c9 --- /dev/null +++ b/tests/test_display_colors.py @@ -0,0 +1,74 @@ +"""Display colours follow AnnData's existing uns convention.""" + +from __future__ import annotations + +import anndata as ad +import numpy as np +import pandas as pd +import scipy.sparse as sp + +from spatialdata_io.experimental.display_colors import ( + GENE_COLORS_KEY, + add_cluster_colors, + add_gene_colors, + palette, +) + + +def _table(n_genes: int = 4) -> ad.AnnData: + return ad.AnnData(X=sp.csr_matrix(np.zeros((3, n_genes), dtype=np.float32))) + + +def test_gene_colours_go_to_uns_not_var() -> None: + # AnnData already has a place for colours: a `_colors` list in uns aligned to an + # ordering. Adding a var column instead would be a new convention for no gain. + table = _table() + assert add_gene_colors(table) == GENE_COLORS_KEY + assert GENE_COLORS_KEY in table.uns + assert "color" not in table.var + assert len(table.uns[GENE_COLORS_KEY]) == table.n_vars + + +def test_colours_are_valid_hex_and_distinct() -> None: + colors = palette(24) + assert all(c.startswith("#") and len(c) == 7 for c in colors) + assert all(all(ch in "0123456789abcdef" for ch in c[1:]) for c in colors) + # Golden-ratio hue stepping should not repeat this early. + assert len(set(colors)) == 24 + + +def test_the_palette_is_deterministic() -> None: + # A gene must keep its colour between runs, or every rebuild reshuffles the display. + assert palette(10) == palette(10) + assert palette(10) == palette(20)[:10] + + +def test_existing_colours_are_not_clobbered() -> None: + table = _table() + table.uns[GENE_COLORS_KEY] = ["#000000"] * table.n_vars + assert add_gene_colors(table) is None + assert table.uns[GENE_COLORS_KEY][0] == "#000000" + assert add_gene_colors(table, overwrite=True) == GENE_COLORS_KEY + assert table.uns[GENE_COLORS_KEY][0] != "#000000" + + +def test_cluster_colours_follow_the_scanpy_key() -> None: + table = _table() + table.obs["cell_type"] = pd.Categorical(["a", "b", "a"]) + assert add_cluster_colors(table, "cell_type") == "cell_type_colors" + assert len(table.uns["cell_type_colors"]) == 2 + + +def test_a_non_categorical_column_is_skipped() -> None: + table = _table() + table.obs["counts"] = [1, 2, 3] + assert add_cluster_colors(table, "counts") is None + assert add_cluster_colors(table, "absent") is None + + +def test_colours_round_trip_through_zarr(tmp_path) -> None: + table = _table() + add_gene_colors(table) + table.write_zarr(tmp_path / "t.zarr") + reloaded = ad.read_zarr(tmp_path / "t.zarr") + assert list(reloaded.uns[GENE_COLORS_KEY]) == list(table.uns[GENE_COLORS_KEY]) diff --git a/tests/test_expression_index.py b/tests/test_expression_index.py new file mode 100644 index 00000000..34236a4e --- /dev/null +++ b/tests/test_expression_index.py @@ -0,0 +1,155 @@ +"""Gene-major expression access without downloading the whole matrix.""" + +from __future__ import annotations + +import anndata as ad +import numpy as np +import pytest +import scipy.sparse as sp + +from spatialdata_io.experimental.expression_index import ( + CSC_LAYER, + GENE_STAT_COLUMNS, + add_csc_layer, + add_gene_statistics, + csc_chunk_size, + write_csc_layer, +) + + +def _table() -> ad.AnnData: + # 4 cells x 3 genes + # gene 0: [1, 0, 0, 3] + # gene 1: [0, 2, 0, 4] + # gene 2: [5, 0, 0, 0] + dense = np.array( + [[1.0, 0.0, 5.0], [0.0, 2.0, 0.0], [0.0, 0.0, 0.0], [3.0, 4.0, 0.0]], + dtype=np.float32, + ) + return ad.AnnData(X=sp.csr_matrix(dense)) + + +def test_statistics_count_the_zero_cells() -> None: + table = _table() + add_gene_statistics(table) + + # gene 0 is [1, 0, 0, 3]: mean 1, E[x^2] = 2.5, var 1.5, 2 of 4 cells non-zero + assert table.var["mean"].iloc[0] == pytest.approx(1.0) + assert table.var["std"].iloc[0] == pytest.approx(np.sqrt(1.5)) + assert table.var["max"].iloc[0] == pytest.approx(3.0) + assert table.var["non_zero"].iloc[0] == pytest.approx(0.5) + + # a gene present in a single cell still reports that cell's value as the max + assert table.var["max"].iloc[2] == pytest.approx(5.0) + assert table.var["non_zero"].iloc[2] == pytest.approx(0.25) + + +def test_statistics_match_a_dense_computation() -> None: + rng = np.random.default_rng(0) + dense = (rng.random((200, 17)) < 0.3) * rng.random((200, 17)) * 10 + table = ad.AnnData(X=sp.csr_matrix(dense.astype(np.float32))) + add_gene_statistics(table) + + assert np.allclose(table.var["mean"], dense.mean(axis=0), atol=1e-5) + assert np.allclose(table.var["std"], dense.std(axis=0), atol=1e-5) + assert np.allclose(table.var["max"], dense.max(axis=0), atol=1e-5) + assert np.allclose(table.var["non_zero"], (dense > 0).mean(axis=0), atol=1e-9) + + +def test_every_documented_column_is_written() -> None: + table = _table() + written = add_gene_statistics(table) + assert written == list(GENE_STAT_COLUMNS) + for column in GENE_STAT_COLUMNS: + assert column in table.var + + +def test_csc_layer_holds_the_same_values_gene_major() -> None: + table = _table() + assert add_csc_layer(table) == CSC_LAYER + + csc = table.layers[CSC_LAYER] + assert sp.isspmatrix_csc(csc) + assert (csc.toarray() == table.X.toarray()).all() + + # The point of the layer: one gene is a contiguous slice, so a client reads two + # chunks instead of the entire matrix. + start, end = csc.indptr[0], csc.indptr[1] + assert sorted(csc.data[start:end]) == [1.0, 3.0] + assert sorted(csc.indices[start:end]) == [0, 3] + + +def test_a_dense_matrix_is_left_alone() -> None: + # Dense is already addressable by column, so a copy would double the storage for + # nothing. + table = ad.AnnData(X=np.zeros((3, 2), dtype=np.float32)) + assert add_csc_layer(table) is None + assert CSC_LAYER not in table.layers + + +def test_no_expression_is_not_an_error() -> None: + table = ad.AnnData(obs=None, var=None, shape=(0, 0)) + assert add_gene_statistics(table) == [] + assert add_csc_layer(table) is None + + +def test_the_layer_round_trips_through_zarr(tmp_path) -> None: + # The whole approach depends on the transpose surviving a write, since a client reads + # it back out of the store rather than from memory. + table = _table() + add_csc_layer(table) + add_gene_statistics(table) + + path = tmp_path / "t.zarr" + table.write_zarr(path) + reloaded = ad.read_zarr(path) + + assert sp.isspmatrix_csc(reloaded.layers[CSC_LAYER]) + assert (reloaded.layers[CSC_LAYER].toarray() == table.X.toarray()).all() + assert reloaded.var["mean"].iloc[0] == pytest.approx(1.0) + + +def test_chunks_are_sized_by_the_average_gene() -> None: + # AnnData chunks for whole-matrix reads: 162,948 non-zeros for Xenium pancreas against + # ~6,915 for one gene, so a gene costs 24x what it needs. These are sized by the gene. + assert csc_chunk_size(nnz=2_607_168, n_genes=377) == pytest.approx(13_830, rel=0.01) + # ...but never so small that the chunk count explodes on a sparse, gene-rich matrix. + assert csc_chunk_size(nnz=1000, n_genes=5000) == 1024 + assert csc_chunk_size(nnz=0, n_genes=0) == 1 + + +def test_written_layer_is_chunked_for_single_gene_reads(tmp_path) -> None: + rng = np.random.default_rng(0) + dense = (rng.random((500, 40)) < 0.2) * rng.random((500, 40)) + table = ad.AnnData(X=sp.csr_matrix(dense.astype(np.float32))) + csc = table.X.tocsc() + + path = tmp_path / "t.zarr" + table.write_zarr(path) + info = write_csc_layer(path, csc) + + import zarr + + group = zarr.open_group(str(path / "layers" / info["layer"]), mode="r") + assert group.attrs["encoding-type"] == "csc_matrix" + assert group.attrs["shape"] == list(csc.shape) + + # A gene must not span many chunks, or reading one costs several requests. + chunk = group["data"].chunks[0] + per_gene = info["nnz"] / info["genes"] + assert chunk >= per_gene + + # indptr is always read whole, so it stays a single chunk. + assert group["indptr"].chunks[0] == group["indptr"].shape[0] + + +def test_written_layer_reads_back_as_the_same_matrix(tmp_path) -> None: + table = _table() + csc = table.X.tocsc() + path = tmp_path / "t.zarr" + table.write_zarr(path) + write_csc_layer(path, csc) + + reloaded = ad.read_zarr(path) + assert sp.isspmatrix_csc(reloaded.layers[CSC_LAYER]) + assert (reloaded.layers[CSC_LAYER].toarray() == table.X.toarray()).all() diff --git a/tests/test_points_parquet.py b/tests/test_points_parquet.py new file mode 100644 index 00000000..969f164a --- /dev/null +++ b/tests/test_points_parquet.py @@ -0,0 +1,435 @@ +"""Tests for the regular-grid Points rewrite. + +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 + +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, + DisplayTransform, + 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 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( + { + "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 + from spatialdata.transformations import Scale, set_transformation + + 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]: + """The canonical element: tile-ordered, carrying only its own columns.""" + out = tmp_path / "points.parquet" + 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 --------------------------------------------------------- + + +@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 + + +@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 = 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): + got = _row_group(directory, manifest, tile_id).num_rows + assert got == expected.get(tile_id, 0), f"tile {tile_id}" + + +@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"]) + md = pq.ParquetFile(directory / manifest["files"][file_index]).metadata + assert md.row_group(local).num_rows == 0 + + +def test_dense_tile_stays_one_physical_row_group(tmp_path: Path) -> None: + """PyArrow otherwise splits a table just above 1 Mi rows despite one write call.""" + n = 1_048_577 + points = pd.DataFrame( + { + "x": np.zeros(n, dtype=np.float32), + "y": np.zeros(n, dtype=np.float32), + "feature_name": pd.Categorical.from_codes(np.zeros(n, dtype=np.int8), ["GENEA"]), + } + ) + out = tmp_path / "dense" + manifest = write_points_regular_grid( + points, + out, + catalog=FeatureCatalog(("GENEA",), 1), + grid=RegularGrid(0, 0, 10, 1, 1), + display_transform=DisplayTransform(((1, 0, 0), (0, 1, 0)), "global"), + render_only=True, + ) + parquet = pq.ParquetFile(out / manifest["files"][0]) + assert parquet.metadata.num_row_groups == 1 + assert parquet.metadata.row_group(0).num_rows == n + + +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 = rendered + rg = _row_group(directory, manifest, 3) + assert [tuple(v) for v in rg[POSITION_COLUMN].to_pylist()] == [(10, 2)] + + +# -- 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: + 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=col) + 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_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 + 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() + + +# -- the render file ---------------------------------------------------------- + + +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_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.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: + """The flat child buffer must be [x0,y0,x1,y1,...] so deck.gl can consume it directly.""" + 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()] + 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(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 _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(), strict=True) + } + + +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 + + +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_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: + 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.""" + 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, render_only=True) + 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") + + +# -- streaming ---------------------------------------------------------------- + + +def _as_partitioned(frame: pd.DataFrame, npartitions: int = 3): + import dask.dataframe as dd + from spatialdata.transformations import Scale + + 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 + + +@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. + """ + 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( + _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) + for tile_id in range(GRID.num_tiles): + 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: + original = points.compute() if hasattr(points, "compute") else points + out = tmp_path / "s.parquet" + 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) + + +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) + + +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_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) diff --git a/tests/test_shapes_parquet.py b/tests/test_shapes_parquet.py new file mode 100644 index 00000000..35a24c8a --- /dev/null +++ b/tests/test_shapes_parquet.py @@ -0,0 +1,245 @@ +"""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 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]: + """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 ------------------------------------------------------------------- + + +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(rendered: tuple[Path, dict]) -> None: + """polygon -> rings -> interleaved uint32 pairs, as get_polygon_data.js walks it.""" + 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.float32(), 2) # interleaved vertices + + +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, _ = rendered + 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(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(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() + names = list(SHAPES_SPEC) + 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 + + +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, 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} + + +# -- 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) + + +# -- cell metadata ------------------------------------------------------------ + + +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 new file mode 100644 index 00000000..52096ff7 --- /dev/null +++ b/tests/test_tiled_access.py @@ -0,0 +1,534 @@ +"""Tests for the profile manifest and the opt-in tiling entry points.""" + +from __future__ import annotations + +import json +import os +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, + ROOT_MANIFEST_KEY, + build_manifest, + validate_manifest, + write_manifest, +) +from spatialdata_io.experimental.regular_grid import RegularGrid +from spatialdata_io.experimental.tiled_access import add_spatial_tiling, xenium_spatially_tiled + +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), strict=True) + ], + index=cells, + ) + shapes = ShapesModel.parse(gdf) + set_transformation(shapes, Scale([2.0, 2.0], axes=("x", "y")), "global") + + # Deliberately reverse instance IDs relative to obs_names: cell_code must follow the + # table's SpatialData instance_key relationship, not happen to match obs_names. + obs = pd.DataFrame( + {"region": pd.Categorical(["cell_boundaries"] * 12), "instance_id": cells[::-1]}, + 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() + # Gene and cell metadata are read from the store's own table, so the profile holds + # only what Zarr cannot serve: the tiled display Parquets. + assert not (profile / "meta_gene.parquet").exists() + assert not (profile / "cell_metadata.parquet").exists() + assert not (profile / "cbg").exists() + assert manifest["spatialdata"]["native"] == ["metadata", "cbg", "images"] + assert manifest["spatialdata"]["store_url"] == "../.." + validate_manifest(manifest, base_path=profile) + + # Shapes are indexed in their own order, while instance_id is reversed in obs. + # Codes must be absolute table-row positions resolved through instance_key. + import spatialdata + + shapes = spatialdata.read_zarr(store).shapes["cell_boundaries"] + cells = [f"cell-{i}" for i in range(12)] + for cell in cells: + assert shapes.loc[cell, "cell_code"] == 11 - int(cell.removeprefix("cell-")) + + +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 + # 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_canonical_layout_uses_store_elements_and_root_manifest(store: Path) -> None: + """Canonical mode adds tiling in place without creating display-only files.""" + import spatialdata + import zarr + + manifest = add_spatial_tiling(store, tile_size_px=10.0, profile_layout="canonical") + + assert not (store / "visualization").exists() + assert manifest["spatialdata"]["store_url"] == "." + assert "native" not in manifest["spatialdata"] + for celldega_key in ( + "technology", + "use_row_groups", + "use_int_index", + "segmentation_approach", + "image_info", + "image_format", + ): + assert celldega_key not in manifest + assert "images" not in manifest["row_group_files"] + assert manifest["source"]["technology"] == "Xenium" + + root = zarr.open_group(store, mode="r") + assert root.attrs[ROOT_MANIFEST_KEY] == manifest + + transcripts = manifest["row_group_files"]["transcripts"] + assert transcripts["directory"] == "points/transcripts/points.parquet" + assert transcripts["position_encoding"] == "separate_columns" + assert transcripts["position_columns"] == ["x", "y"] + assert transcripts["feature_column"] == "feature_name" + assert transcripts["feature_encoding"] == "dictionary" + + cells = manifest["row_group_files"]["cell_segmentation"] + assert cells["path"] == "shapes/cell_boundaries/shapes.parquet" + assert "directory" not in cells + assert cells["geometry_column"] == "geometry" + assert cells["geometry_encoding"] == "geoarrow.polygon" + assert cells["geometry_is_lossy"] is False + assert cells["display_transform"]["affine_matrix"][0][0] == 2.0 + + validate_manifest(manifest, base_path=store) + reopened = spatialdata.read_zarr(store) + assert len(reopened.points["transcripts"]) == 200 + assert len(reopened.shapes["cell_boundaries"]) == 12 + + +def test_invalid_profile_layout_is_rejected_before_writing(store: Path) -> None: + with pytest.raises(ValueError, match="profile_layout must be"): + add_spatial_tiling(store, profile_layout="unknown") + + +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_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 "display_xy" not in cols + assert "feature_code" not in cols + + +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") + + +# -- 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_one_shot_indexes_table_before_initial_write( + store: Path, tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """The opt-in Xenium path must not write, delete, then rewrite its table.""" + import spatialdata + + source = spatialdata.read_zarr(store) + monkeypatch.setattr("spatialdata_io.readers.xenium.xenium", lambda *_args, **_kwargs: source) + + def fail_on_delete(*_args, **_kwargs): + raise AssertionError("the one-shot path must not delete the freshly written table") + + monkeypatch.setattr(spatialdata.SpatialData, "delete_element_from_disk", fail_on_delete) + + out = tmp_path / "one-shot.zarr" + manifest = xenium_spatially_tiled(tmp_path / "raw-placeholder", out, tile_size_px=10.0) + table = spatialdata.read_zarr(out).tables["table"] + assert {"mean", "std", "max", "non_zero"} <= set(table.var.columns) + assert "gene_colors" in table.uns + assert "X_csc" in table.layers + assert manifest["spatialdata"]["expression_index"]["csc"]["layer"] == "X_csc" + + +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"]) + + +# -- 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": cells}, 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. + # The manifest has to say so: a client reading gene names from `var` alone would + # otherwise have no entry for Blank-1's feature_code and would drop its colour. + features = manifest["feature_catalog"] + assert features["n_genes"] == 3 + assert features["extra_features"] == ["Blank-1"] + + # 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 + + +def test_no_table_manifest_does_not_advertise_table_components(store: Path) -> None: + manifest = add_spatial_tiling( + store, + table_element=None, + shapes_element=None, + tile_size_px=10.0, + ) + assert manifest["spatialdata"]["native"] == ["images"] + assert "table" not in manifest["spatialdata"] + assert manifest["feature_catalog"]["names"] == [ + "GENEA", + "GENEB", + "GENEC", + "NegControlProbe_0001", + ]