diff --git a/src/spatialdata_io/experimental/cbg_parquet.py b/src/spatialdata_io/experimental/cbg_parquet.py new file mode 100644 index 00000000..92d8096b --- /dev/null +++ b/src/spatialdata_io/experimental/cbg_parquet.py @@ -0,0 +1,181 @@ +"""Write a gene-major cell-by-gene matrix as one row group per gene. + +This is what lets the overview render cheaply: selecting a gene fetches a single row +group holding that gene's non-zero cells, instead of touching any transcript data. + +The schema matches Celldega's existing CBG reader exactly -- columns ``cell_id`` +(the integer cell code, not a string barcode), ``expression``, ``gene``, plus +``gene_to_row_group`` / ``num_genes`` in the Parquet schema metadata. + +One deliberate difference from Celldega's own writer: every gene in the catalog gets a +row group, in catalog order, even if it has no non-zero cells. That preserves the +invariant ``feature_code == cbg row group``, so a transcript's feature code and its +expression vector are addressed by the same integer. ``gene_to_row_group`` is still +written, so a client that looks the mapping up rather than assuming it works unchanged. +""" + +from __future__ import annotations + +import json +import shutil +from pathlib import Path +from typing import Any + +import numpy as np +import pyarrow as pa +import pyarrow.parquet as pq +import scipy.sparse as sp + +from spatialdata_io.experimental.feature_catalog import FeatureCatalog +from spatialdata_io.experimental.regular_grid import DEFAULT_MAX_ROW_GROUPS_PER_FILE + +__all__ = ["write_cbg_row_groups"] + + +def write_cbg_row_groups( + table: Any, + output_dir: str | Path, + *, + catalog: FeatureCatalog, + cell_codes: dict[str, int] | None = None, + layer: str | None = None, + max_row_groups_per_file: int = DEFAULT_MAX_ROW_GROUPS_PER_FILE, + compression: str = "zstd", + overwrite: bool = False, +) -> dict[str, Any]: + """Write an AnnData table as gene-major CBG row groups. + + Parameters + ---------- + table + The annotating :class:`anndata.AnnData` table. + output_dir + Directory to write the chunk files into. Written atomically. + catalog + Feature catalog; its gene order defines the row-group order. Every gene in the + catalog must be present in ``table.var_names``. + cell_codes + Mapping from ``obs`` name to integer cell code. Defaults to positional order, + which is what :func:`write_shapes_regular_grid` uses by default. + layer + Optional ``table.layers`` key to read instead of ``table.X``. + max_row_groups_per_file + Genes per chunk file. + compression + Parquet compression codec. + overwrite + Replace ``output_dir`` if it exists. + + Returns + ------- + The manifest fragment describing the written files. + """ + output_dir = Path(output_dir) + if output_dir.exists() and not overwrite: + raise FileExistsError(f"{output_dir} exists; pass overwrite=True to replace it") + + var_names = list(table.var_names) + positions = {g: i for i, g in enumerate(var_names)} + missing = [g for g in catalog.genes if g not in positions] + if missing: + raise ValueError( + f"{len(missing)} catalog gene(s) are absent from table.var_names " + f"(e.g. {missing[:3]}); the CBG would not cover every gene code." + ) + + if cell_codes is None: + codes = np.arange(table.n_obs, dtype=np.uint32) + else: + unknown = [k for k in table.obs_names if k not in cell_codes] + if unknown: + raise ValueError(f"{len(unknown)} table cell(s) have no cell_code (e.g. {unknown[:3]})") + codes = np.fromiter((cell_codes[k] for k in table.obs_names), dtype=np.uint32, count=table.n_obs) + + matrix = table.layers[layer] if layer is not None else table.X + # CSC gives O(1) access to a single gene's column, which is the whole point here. + csc = matrix.tocsc() if sp.issparse(matrix) else sp.csc_matrix(np.asarray(matrix)) + csc.sort_indices() + + schema = pa.schema( + [ + pa.field("cell_id", pa.uint32()), + pa.field("expression", pa.float32()), + pa.field("gene", pa.string()), + ] + ) + + gene_to_row_group = {gene: i for i, gene in enumerate(catalog.genes)} + n_genes = len(catalog.genes) + n_files = -(-n_genes // max_row_groups_per_file) + width = len(str(n_files - 1)) if n_files > 1 else 1 + filenames = [f"chunk_{i:0{width}d}.parquet" for i in range(n_files)] + + schema = schema.with_metadata( + { + b"gene_to_row_group": json.dumps(gene_to_row_group).encode(), + b"storage_mode": b"row_groups_cbg_chunked", + b"num_genes": str(n_genes).encode(), + b"max_row_groups_per_file": str(max_row_groups_per_file).encode(), + b"profile": b"grid_files_v1", + } + ) + + staging = output_dir.with_name(output_dir.name + ".tmp") + if staging.exists(): + shutil.rmtree(staging) + staging.mkdir(parents=True) + + n_values = 0 + try: + writer: pq.ParquetWriter | None = None + current_file = -1 + for row_group, gene in enumerate(catalog.genes): + file_index = row_group // max_row_groups_per_file + if file_index != current_file: + if writer is not None: + writer.close() + writer = pq.ParquetWriter( + staging / filenames[file_index], schema, compression=compression, write_statistics=False + ) + current_file = file_index + + col = positions[gene] + start, end = csc.indptr[col], csc.indptr[col + 1] + cells = codes[csc.indices[start:end]] + values = csc.data[start:end] + # Explicit zeros can survive in a sparse matrix; the CBG stores non-zeros only. + keep = values != 0 + cells, values = cells[keep], values[keep] + n_values += len(values) + + assert writer is not None + writer.write_table( + pa.table( + { + "cell_id": pa.array(cells, type=pa.uint32()), + "expression": pa.array(values, type=pa.float32()), + "gene": pa.array([gene] * len(values), type=pa.string()), + }, + schema=schema, + ) + ) + if writer is not None: + writer.close() + except BaseException: + shutil.rmtree(staging, ignore_errors=True) + raise + + if output_dir.exists(): + shutil.rmtree(output_dir) + staging.rename(output_dir) + + return { + "directory": output_dir.name, + "files": filenames, + "max_row_groups_per_file": max_row_groups_per_file, + "total_row_groups": n_genes, + "num_genes": n_genes, + "gene_to_row_group": gene_to_row_group, + "n_values": n_values, + "row_group_equals_feature_code": True, + } diff --git a/src/spatialdata_io/experimental/feature_catalog.py b/src/spatialdata_io/experimental/feature_catalog.py new file mode 100644 index 00000000..81e4b21f --- /dev/null +++ b/src/spatialdata_io/experimental/feature_catalog.py @@ -0,0 +1,269 @@ +"""Stable integer codes for transcript features. + +Rendering a transcript layer needs a compact integer per point, not a string. This module +builds the mapping ``feature name <-> feature_code`` and pins two properties the rest of +the profile depends on: + +1. Genes come first, in the annotating table's ``var_names`` order, so a gene's + ``feature_code`` *is* its row-group index in the cell-by-gene file. No second lookup + table, and no browser-side string join. +2. Non-gene features (negative controls, unassigned codewords) are kept, but are coded + *above* every gene and flagged. They are never silently folded into a real gene, which + would fabricate expression. + +The split point is recorded as ``n_genes`` so a client can tell the two apart from the +manifest alone. +""" + +from __future__ import annotations + +from collections.abc import Iterable, Sequence +from dataclasses import dataclass +from typing import Any + +import numpy as np +import pandas as pd +from numpy.typing import NDArray + +__all__ = ["FeatureCatalog", "CONTROL_PREFIXES"] + +#: Feature-name prefixes that 10x uses for non-gene codewords. Used only for reporting; +#: catalog membership is decided by absence from the table's ``var_names``, not by prefix. +CONTROL_PREFIXES = ( + "NegControlProbe", + "NegControlCodeword", + "UnassignedCodeword", + "antisense", + "BLANK", + "DeprecatedCodeword", + "Intergenic", +) + + +@dataclass(frozen=True) +class FeatureCatalog: + """An ordered feature vocabulary with genes first. + + Parameters + ---------- + names + All feature names. Position in this list is the feature code. + n_genes + Number of leading entries that are genes present in the annotating table. Codes + ``>= n_genes`` are non-gene features. + """ + + names: tuple[str, ...] + n_genes: int + + def __post_init__(self) -> None: + if len(set(self.names)) != len(self.names): + raise ValueError("feature names must be unique") + if not 0 <= self.n_genes <= len(self.names): + raise ValueError(f"n_genes={self.n_genes} out of range for {len(self.names)} features") + + def __len__(self) -> int: + return len(self.names) + + @property + def genes(self) -> tuple[str, ...]: + """The gene names, in table ``var_names`` order.""" + return self.names[: self.n_genes] + + @property + def controls(self) -> tuple[str, ...]: + """The non-gene feature names.""" + return self.names[self.n_genes :] + + @property + def dtype(self) -> np.dtype[Any]: + """Smallest unsigned integer dtype that can hold every code.""" + return np.dtype(np.uint16) if len(self.names) <= np.iinfo(np.uint16).max else np.dtype(np.uint32) + + # -- construction --------------------------------------------------------- + + @classmethod + def from_features_and_table( + cls, + feature_names: Iterable[str], + var_names: Sequence[str], + ) -> FeatureCatalog: + """Build a catalog from the observed transcript features and the table's genes. + + Genes are taken in ``var_names`` order. Any observed feature absent from + ``var_names`` is appended as a control, sorted so the catalog is reproducible. + + Raises + ------ + ValueError + If ``var_names`` contains duplicates, which would make codes ambiguous. + """ + genes = list(var_names) + if len(set(genes)) != len(genes): + raise ValueError("var_names contains duplicates; feature codes would be ambiguous") + + observed = set(feature_names) + gene_set = set(genes) + controls = sorted(observed - gene_set) + return cls(names=tuple(genes) + tuple(controls), n_genes=len(genes)) + + @classmethod + def from_points_and_table( + cls, + points: Any, + table: Any, + feature_key: str = "feature_name", + ) -> FeatureCatalog: + """Build a catalog directly from a Points element and its annotating table. + + Handles the dask categorical read back from a zarr store, whose categories are + lazily "unknown" and raise on access until realized. + """ + col = points[feature_key] + if hasattr(col, "cat"): + try: + features = list(col.cat.categories) + except (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_frame(self, expression: Any | None = None) -> pd.DataFrame: + """Return the catalog as the ``meta_gene.parquet`` table. + + The layout follows Celldega's own ``meta_gene.parquet``, because a client reads it + by convention rather than through the manifest: the **index** is the gene name, + and the columns are ``mean``, ``std``, ``max``, ``non-zero`` and ``color``. A + client with no ``color`` column or no index finds no genes at all, which shows up + as a viewer with no transcript controls rather than as an error. + + ``feature_code`` and ``is_gene`` are carried alongside as profile additions. + + Parameters + ---------- + expression + Optional cell-by-gene matrix (cells x features, in catalog order) used to + compute the per-gene statistics. Without it the statistics are zero, which is + valid but leaves the viewer's gene ranking flat. + """ + import colorsys + + n = len(self.names) + stats = {k: np.zeros(n, dtype=np.float64) for k in ("mean", "std", "max", "non-zero")} + if expression is not None: + stats.update(_expression_stats(expression, n)) + + # A hue sweep, matching the look of Celldega's generated palette. Blank/control + # features are white so they read as "not a gene" in the UI. + colors = [] + for i, name in enumerate(self.names): + if i >= self.n_genes or "Blank" in name: + colors.append("#FFFFFF") + else: + r, g, b = colorsys.hsv_to_rgb(i / max(1, self.n_genes), 0.7, 0.9) + colors.append(f"#{int(r * 255):02x}{int(g * 255):02x}{int(b * 255):02x}") + + frame = pd.DataFrame( + { + **stats, + "color": colors, + "feature_code": np.arange(n, dtype=self.dtype), + "is_gene": np.arange(n) < self.n_genes, + }, + # The index is deliberately unnamed: pandas then writes it as + # '__index_level_0__', which is the only name a client looks for. A named + # index becomes a column of that name instead, and the gene list reads empty. + index=pd.Index(list(self.names)), + ) + # Deliberately NOT sorted. A client builds its integer gene id from a feature's + # *row position* in this file, and colours transcripts by indexing an array built + # the same way. Sorting would break the correspondence with feature_code, so most + # colour lookups would miss and the transcripts would render transparent. + return frame + + def to_manifest_dict(self) -> dict[str, Any]: + """Summary for the profile manifest. The full mapping lives in ``meta_gene.parquet``.""" + return { + "n_features": len(self.names), + "n_genes": self.n_genes, + "feature_code_dtype": self.dtype.name, + "gene_codes_match_cbg_row_groups": True, + } + + +def _expression_stats(matrix: Any, n_features: int) -> dict[str, NDArray[np.float64]]: + """Per-feature mean, std, max and non-zero fraction from a cells x features matrix. + + Computed column-wise on the sparse matrix rather than densifying it, which for a + 5,000-gene panel would otherwise be several GB. + """ + import scipy.sparse as sp + + stats = {k: np.zeros(n_features, dtype=np.float64) for k in ("mean", "std", "max", "non-zero")} + if matrix is None: + return stats + + csc = matrix.tocsc() if sp.issparse(matrix) else sp.csc_matrix(np.asarray(matrix)) + n_cells = csc.shape[0] + if n_cells == 0: + return stats + + for col in range(min(n_features, csc.shape[1])): + values = csc.data[csc.indptr[col] : csc.indptr[col + 1]] + values = values[values != 0].astype(np.float64) + if values.size == 0: + continue + mean = float(values.sum()) / n_cells + # Variance over all cells, counting the implicit zeros. + stats["mean"][col] = mean + stats["std"][col] = float(np.sqrt(max(0.0, (values**2).sum() / n_cells - mean**2))) + stats["max"][col] = float(values.max()) + stats["non-zero"][col] = values.size / n_cells + return stats diff --git a/src/spatialdata_io/experimental/manifest.py b/src/spatialdata_io/experimental/manifest.py new file mode 100644 index 00000000..993a2e21 --- /dev/null +++ b/src/spatialdata_io/experimental/manifest.py @@ -0,0 +1,192 @@ +"""The visualization profile manifest. + +The manifest is what turns a pile of Parquet files into a discoverable profile: it tells a +client the tile geometry, which files hold which row groups, and which columns to project. +Without it a client would have to infer the layout, which is exactly what this profile +exists to avoid. + +Key names follow Celldega's existing ``landscape_parameters.json`` so that its reader can +consume the manifest unchanged (``use_row_groups``, ``tile_grid``, ``row_group_files``, +``technology``, ``image_info``). Profile-specific additions live under ``profile`` and in +each entry's column names, so a client that does not understand them still finds the keys +it expects. +""" + +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any + +from spatialdata_io.experimental.regular_grid import RegularGrid + +__all__ = [ + "PROFILE_NAME", + "PROFILE_VERSION", + "MANIFEST_FILENAME", + "build_manifest", + "validate_manifest", + "write_manifest", +] + +PROFILE_NAME = "grid_files_v1" +PROFILE_VERSION = "0.1.0" +MANIFEST_FILENAME = "landscape_parameters.json" + + +def build_manifest( + *, + grid: RegularGrid, + technology: str = "Xenium", + transcripts: dict[str, Any] | None = None, + cell_segmentation: dict[str, Any] | None = None, + cbg: dict[str, Any] | None = None, + images: dict[str, Any] | None = None, + image_info: list[dict[str, Any]] | None = None, + image_dimensions: dict[str, Any] | None = None, + max_pyramid_zoom: int | None = None, + fixed_path_assets: dict[str, Any] | None = None, + source: dict[str, Any] | None = None, +) -> dict[str, Any]: + """Assemble the profile manifest from the fragments returned by each writer. + + 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 ------------------------- + # Names and shapes follow a DegaFiles landscape_parameters.json so the reader + # needs no special-casing for a SpatialData store. + "technology": technology, + "use_row_groups": True, + "use_int_index": True, + "segmentation_approach": ["default"], + "tile_size": grid.tile_size_px, + "tile_grid": grid.to_manifest_dict(), + "row_group_files": row_group_files, + "image_info": image_info or [], + "image_format": ".webp", + # -- profile identification ------------------------------------------- + "profile": PROFILE_NAME, + "profile_version": PROFILE_VERSION, + } + if 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 ("technology", "use_row_groups", "tile_grid", "row_group_files"): + if key not in manifest: + raise ValueError(f"manifest is missing required key {key!r}") + + grid = RegularGrid.from_manifest_dict(manifest["tile_grid"]) + files = manifest["row_group_files"] + base = Path(base_path) if base_path is not None else None + + for name in ("transcripts", "cell_segmentation"): + entry = files.get(name) + if entry is None: + continue + expected = grid.num_tiles + if entry.get("total_row_groups") != expected: + raise ValueError( + f"{name}: total_row_groups={entry.get('total_row_groups')} does not match the " + f"tile grid ({grid.num_tiles_x} x {grid.num_tiles_y} = {expected} tiles). " + f"The row-group index formula would address the wrong tiles." + ) + _check_paths(name, entry, base) + + cbg = files.get("cbg") + if cbg is not None: + mapping = cbg.get("gene_to_row_group", {}) + if not mapping: + raise ValueError("cbg: gene_to_row_group is empty; no gene could be selected") + if cbg.get("total_row_groups") != len(mapping): + raise ValueError( + f"cbg: total_row_groups={cbg.get('total_row_groups')} does not match " + f"{len(mapping)} entries in gene_to_row_group" + ) + max_rg = cbg.get("max_row_groups_per_file", 1) + n_files = len(cbg.get("files", [])) + if mapping and max(mapping.values()) >= n_files * max_rg: + raise ValueError( + f"cbg: gene_to_row_group references row group {max(mapping.values())} but only " + f"{n_files} file(s) x {max_rg} row groups are listed" + ) + _check_paths("cbg", cbg, base) + + +def _check_paths(name: str, entry: dict[str, Any], base: Path | None) -> None: + """Verify the declared files exist, and that the entry names them coherently.""" + if "files" in entry: + if "directory" not in entry: + raise ValueError(f"{name}: entry lists 'files' but no 'directory'") + n_files = len(entry["files"]) + expected = -(-entry["total_row_groups"] // entry["max_row_groups_per_file"]) + if n_files != expected: + raise ValueError( + f"{name}: lists {n_files} file(s) but {entry['total_row_groups']} row groups at " + f"{entry['max_row_groups_per_file']} per file need {expected}" + ) + if base is not None: + for f in entry["files"]: + if not (base / entry["directory"] / f).exists(): + raise ValueError(f"{name}: declared file {entry['directory']}/{f} does not exist") + elif "path" in entry: + if base is not None and not (base / entry["path"]).exists(): + raise ValueError(f"{name}: declared path {entry['path']} does not exist") + else: + raise ValueError(f"{name}: entry has neither 'files' nor 'path'") + + +def write_manifest(manifest: dict[str, Any], directory: str | Path) -> Path: + """Write the manifest as ``landscape_parameters.json`` and return its path.""" + directory = Path(directory) + directory.mkdir(parents=True, exist_ok=True) + path = directory / MANIFEST_FILENAME + path.write_text(json.dumps(manifest, indent=2, sort_keys=False)) + return path diff --git a/src/spatialdata_io/experimental/points_parquet.py b/src/spatialdata_io/experimental/points_parquet.py new file mode 100644 index 00000000..55c880ca --- /dev/null +++ b/src/spatialdata_io/experimental/points_parquet.py @@ -0,0 +1,559 @@ +"""Rewrite a Points element into regular-grid row groups. + +The output keeps every canonical column and adds two render-oriented ones: + +``display_xy`` + ``fixed_size_list[2]`` of integer level-0 pixel coordinates. The Arrow child + buffer is therefore already ``[x0, y0, x1, y1, ...]`` -- directly usable as a deck.gl + binary ``getPosition`` attribute with no interleaving step in the browser. +``feature_code`` + Small unsigned integer into the :class:`FeatureCatalog`. + +Physical row order changes (rows are grouped by tile), but no row is added, dropped or +altered, and the DataFrame index is preserved so the reordering is fully traceable. + +Row groups are written one-per-logical-tile *including empty tiles*, so a client can find +a tile's data from the tile formula alone, with no lookup table and no reliance on Parquet +statistics. Output is split across several files because a reader must fetch a whole +footer before reading any row group, and footer size grows with row-group count. +""" + +from __future__ import annotations + +import json +import shutil +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import numpy as np +import pandas as pd +import pyarrow as pa +import pyarrow.parquet as pq +from numpy.typing import NDArray + +from spatialdata_io.experimental.feature_catalog import FeatureCatalog +from spatialdata_io.experimental.regular_grid import ( + DEFAULT_MAX_ROW_GROUPS_PER_FILE, + RegularGrid, +) + +__all__ = [ + "POSITION_COLUMN", + "FEATURE_COLUMN", + "DisplayTransform", + "write_points_regular_grid", +] + +#: Column holding interleaved integer pixel positions. +POSITION_COLUMN = "display_xy" +#: Column holding the integer feature code. +FEATURE_COLUMN = "feature_code" + +#: Largest coordinate float32 represents exactly to better than 0.01 px. +_FLOAT32_SAFE_MAX = 2**24 + + +@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, +) -> 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 -- and parquet-wasm's projection is broken + # anyway (any `columns` argument corrupts the IPC stream it emits). + table = pa.table( + { + POSITION_COLUMN: _interleaved_positions(px, py), + FEATURE_COLUMN: pa.array(catalog.encode(df[feature_key])), + } + ) + return table, tile_ids + + # The canonical element keeps only its own columns. The render columns live in a + # separate file, for two reasons: a nested Arrow column cannot survive dask's parquet + # round-trip (SpatialData.write() either fails or silently returns it as a string), + # and a standalone render file means a viewer reads every column of it, so no column + # projection is needed -- which matters because parquet-wasm's projection is broken. + # + # Any render columns left by an earlier version are dropped, so re-tiling a store + # written before this change cleans it up rather than preserving them. + stale = [c for c in (POSITION_COLUMN, FEATURE_COLUMN, _TILE_ID) if c in df.columns] + if df.attrs or stale or categories is not None: + df = df.copy(deep=False) + # 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) + + return pa.Table.from_pandas(df, preserve_index=True), tile_ids + + +def _sorted_by_tile(table: pa.Table, tile_ids: NDArray[np.int64]) -> tuple[pa.Table, NDArray[np.int64]]: + """Group rows by tile. The sort is stable, so the rewrite is deterministic.""" + order = np.argsort(tile_ids, kind="stable") + return table.take(pa.array(order)), tile_ids[order] + + +def _write_tile_row_groups( + writer: pq.ParquetWriter, + table: pa.Table, + sorted_tile_ids: NDArray[np.int64], + tile_range: range, + schema: pa.Schema, +) -> None: + """Write one row group per tile in ``tile_range``, empty tiles included. + + Empty tiles must still occupy a row group, since that is what makes + ``row_group_index == tile_id`` hold without a lookup table. + """ + offsets = np.searchsorted(sorted_tile_ids, np.array([*tile_range, tile_range.stop]), side="left") + for i in range(len(tile_range)): + start, end = int(offsets[i]), int(offsets[i + 1]) + writer.write_table(table.slice(start, end - start) if end > start else schema.empty_table()) + + +def _iter_chunks(points: Any) -> Any: + """Yield the element one partition at a time, or once if it is already in memory.""" + if hasattr(points, "npartitions"): + for i in range(points.npartitions): + yield points.partitions[i].compute() + else: + yield points + + +def _known_categories(points: Any, feature_key: str) -> Any | None: + """Return the element's full category list, so every chunk shares one dictionary.""" + col = points[feature_key] + if not hasattr(col, "cat"): + return None + try: + return list(col.cat.categories) + except (NotImplementedError, AttributeError): + return list(col.cat.as_known().cat.categories) + + +def write_points_regular_grid( + points: Any, + output_dir: str | Path, + *, + 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, + 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, + ) + + 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, + ) + 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) + 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, +) -> 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, + ) + 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) + 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..486757bc --- /dev/null +++ b/src/spatialdata_io/experimental/shapes_parquet.py @@ -0,0 +1,326 @@ +"""Rewrite a Shapes element into regular-grid row groups. + +Adds two render-oriented columns beside the canonical geometry: + +``display_geometry`` + ``list[2]>>`` -- polygon -> rings -> interleaved integer + pixel vertices. The nesting is chosen so a client can lift deck.gl's ``getPolygon`` + straight out of the flat coordinate child buffer and ``startIndices`` out of the list + offsets, with no WKB parsing and no per-vertex JavaScript objects. +``cell_code`` + Positional index into the annotating table, so cells can be coloured from a + cell-by-gene vector without a string join in the browser. + +``display_geometry`` is explicitly a *lossy display* representation: only the exterior +ring is kept, and for a MultiPolygon only its largest part. The canonical geometry column +is written through unchanged, and the GeoParquet metadata is preserved so the file is +still readable by :func:`geopandas.read_parquet` and by SpatialData itself. + +Each cell is assigned to exactly one tile, by its centroid in display pixel space. A +polygon whose outline crosses into a neighbouring tile is *not* duplicated -- duplicating +would inflate the file and make cell counts wrong. +""" + +from __future__ import annotations + +import json +import shutil +from pathlib import Path +from typing import Any + +import numpy as np +import pyarrow as pa +import pyarrow.parquet as pq +import shapely +from numpy.typing import NDArray + +from spatialdata_io.experimental.points_parquet import DisplayTransform, _to_display_pixels +from spatialdata_io.experimental.regular_grid import ( + DEFAULT_MAX_ROW_GROUPS_PER_FILE, + RegularGrid, +) + +__all__ = [ + "GEOMETRY_COLUMN", + "CELL_CODE_COLUMN", + "write_shapes_regular_grid", + "write_cell_metadata", +] + +#: Column holding the nested integer-pixel display polygons. +GEOMETRY_COLUMN = "display_geometry" +#: Column holding the positional cell index. +CELL_CODE_COLUMN = "cell_code" + + +def _largest_polygon(geom: Any) -> Any: + """Reduce a MultiPolygon to its largest part, matching Celldega's existing behaviour.""" + if geom is None or geom.is_empty: + return geom + if geom.geom_type == "MultiPolygon": + return max(geom.geoms, key=lambda g: g.area) + return geom + + +def _exterior_only(geometry: Any) -> NDArray[Any]: + """Return an array of single-ring polygons: largest part, exterior ring only.""" + geoms = np.asarray([_largest_polygon(g) for g in geometry], dtype=object) + rings = shapely.get_exterior_ring(geoms) + return shapely.polygons(rings) + + +def _display_geometry_array( + geometry: Any, transform: DisplayTransform +) -> tuple[pa.ListArray, NDArray[np.uint32], NDArray[np.uint32]]: + """Build the nested display-geometry array and the per-cell display centroids. + + Returns the Arrow array plus the integer pixel centroid coordinates used for tiling. + """ + simple = _exterior_only(geometry) + + # to_ragged_array gives exactly the buffers the nested Arrow layout needs: + # a flat (N, 2) coordinate array plus ring and polygon offsets. + _, coords, offsets = shapely.to_ragged_array(simple) + ring_offsets, polygon_offsets = offsets + + px, py = _to_display_pixels(coords[:, 0], coords[:, 1], transform) + flat = np.empty(px.size * 2, dtype=np.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) -> pa.Table: + """Convert a GeoDataFrame to Arrow while keeping the GeoParquet ``geo`` metadata. + + ``GeoDataFrame.to_arrow()`` emits GeoArrow *extension* metadata but not the GeoParquet + ``geo`` schema key, which is added by ``to_parquet()``. Without that key the rewritten + file is no longer readable by :func:`geopandas.read_parquet` or by SpatialData. We + therefore reuse the same conversion ``to_parquet`` performs, rather than writing the + file once just to recover its metadata. + """ + try: + from geopandas.io.arrow import _geopandas_to_arrow + except ImportError as exc: # pragma: no cover - depends on geopandas internals + raise RuntimeError( + "could not access geopandas' arrow conversion; a geopandas version with " + "geopandas.io.arrow._geopandas_to_arrow is required to preserve GeoParquet metadata" + ) from exc + + table = _geopandas_to_arrow(shapes, index=None, geometry_encoding="WKB") + if b"geo" not in (table.schema.metadata or {}): # pragma: no cover - defensive + raise RuntimeError("geopandas did not produce GeoParquet 'geo' metadata") + return table + + +def write_shapes_regular_grid( + shapes: Any, + output_path: str | Path, + *, + grid: RegularGrid, + display_transform: DisplayTransform | None = None, + coordinate_system: str = "global", + cell_index: Any | None = None, + max_row_groups_per_file: int = DEFAULT_MAX_ROW_GROUPS_PER_FILE, + compression: str = "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 + Index defining ``cell_code`` order, normally the annotating table's ``obs_names``. + Defaults to the shapes' own index order. + max_row_groups_per_file + Row groups per chunk file. + compression + Parquet compression codec. + overwrite + Replace an existing output. + + Returns + ------- + The manifest fragment describing the written file(s). + """ + output_path = Path(output_path) + if output_path.exists() and not overwrite: + raise FileExistsError(f"{output_path} exists; pass overwrite=True to replace it") + + transform = display_transform or DisplayTransform.from_element(shapes, coordinate_system) + display, cx, cy = _display_geometry_array(shapes.geometry, transform) + + # 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) + + if cell_index is None: + codes = np.arange(len(shapes), dtype=np.uint32) + else: + positions = {k: i for i, k in enumerate(cell_index)} + missing = [k for k in shapes.index if k not in positions] + if missing: + raise ValueError( + f"{len(missing)} shape(s) are absent from cell_index (e.g. {missing[:3]}); " + f"cell_code would be undefined. Pass the table's obs_names for these cells." + ) + codes = np.fromiter((positions[k] for k in shapes.index), dtype=np.uint32, count=len(shapes)) + + 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 + writer.write_table(table.slice(start, end - start) if end > start else schema.empty_table()) + if writer is not None: + writer.close() + except BaseException: + if staging.exists(): + shutil.rmtree(staging) if staging.is_dir() else staging.unlink() + raise + + if output_path.exists(): + shutil.rmtree(output_path) if output_path.is_dir() else output_path.unlink() + staging.rename(output_path) + + fragment: dict[str, Any] = { + "geometry_column": GEOMETRY_COLUMN, + "cell_id_column": CELL_CODE_COLUMN, + "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": True, + "geometry_note": "exterior ring of the largest polygon part; canonical geometry retained", + "tile_grid": grid.to_manifest_dict(), + } + if single_file: + fragment["path"] = output_path.name + else: + fragment["directory"] = output_path.name + fragment["files"] = filenames + return fragment + + +def write_cell_metadata( + shapes: Any, + output_path: str | Path, + *, + display_transform: DisplayTransform | None = None, + coordinate_system: str = "global", + cell_index: Any | None = None, +) -> dict[str, Any]: + """Write the per-cell centroid table used for the overview scatter layer. + + Cells, not sampled transcripts, are the overview representation, so a client needs + every cell's centroid up front. This is a single small file (a few MB for ~10^5 cells) + rather than a tiled one, because the whole set is wanted at once. + + Row order is significant: a client takes a cell's integer id from its *position* here, + so the order must match ``cell_index`` (the annotating table's ``obs`` order) and hence + the ``cell_code`` written into the tiled shapes and the CBG. + + The schema matches Celldega's ``cell_metadata.parquet``: ``name`` plus ``geometry`` as + ``list`` holding ``[x, y]`` in display pixels. + """ + output_path = Path(output_path) + transform = display_transform or DisplayTransform.from_element(shapes, coordinate_system) + + simple = _exterior_only(shapes.geometry) + centroids = shapely.centroid(simple) + cx, cy = _to_display_pixels(shapely.get_x(centroids), shapely.get_y(centroids), transform) + + order = np.arange(len(shapes)) + names = [str(k) for k in shapes.index] + if cell_index is not None: + position = {k: i for i, k in enumerate(shapes.index)} + missing = [k for k in cell_index if k not in position] + if missing: + raise ValueError( + f"{len(missing)} cell(s) in cell_index have no shape (e.g. {missing[:3]}); " + f"their centroid would be undefined." + ) + order = np.fromiter((position[k] for k in cell_index), dtype=np.int64, count=len(cell_index)) + names = [str(k) for k in cell_index] + + flat = np.empty(len(order) * 2, dtype=np.float64) + flat[0::2] = cx[order] + flat[1::2] = cy[order] + geometry = pa.ListArray.from_arrays(pa.array(np.arange(len(order) + 1, dtype=np.int32) * 2), pa.array(flat)) + + table = pa.table({"name": pa.array(names, pa.string()), "geometry": geometry}) + output_path.parent.mkdir(parents=True, exist_ok=True) + pq.write_table(table, output_path, compression="zstd") + + return { + "path": output_path.name, + "n_cells": int(table.num_rows), + "position_encoding": "list", + "coordinate_space": "image-pixel", + "order": "annotating table obs order (position == cell_code)", + } diff --git a/src/spatialdata_io/experimental/tiled_access.py b/src/spatialdata_io/experimental/tiled_access.py new file mode 100644 index 00000000..ef5b2bd8 --- /dev/null +++ b/src/spatialdata_io/experimental/tiled_access.py @@ -0,0 +1,449 @@ +"""Opt-in spatial tiling for a SpatialData store. + +Two entry points, both a single call: + +:func:`add_spatial_tiling` + Add the profile to a store that already exists. +:func:`xenium_spatially_tiled` + Read raw Xenium and write a tiled store in one go. + +The one-shot path is internally ``read -> write -> tile``, because the tiling rewrites +*written* Parquet. + +Everything here is additive and opt-in. A store that has been tiled is still an ordinary +SpatialData store: :func:`spatialdata.read_zarr` works unchanged, the canonical columns and +geometries are untouched, and a client that does not know about the profile simply ignores +the extra columns and the manifest. +""" + +from __future__ import annotations + +import shutil +from pathlib import Path +from typing import Any + +import numpy as np + +from spatialdata_io.experimental.cbg_parquet import write_cbg_row_groups +from spatialdata_io.experimental.feature_catalog import FeatureCatalog +from spatialdata_io.experimental.manifest import ( + PROFILE_NAME, + build_manifest, + validate_manifest, + write_manifest, +) +from spatialdata_io.experimental.points_parquet import ( + DisplayTransform, + write_points_regular_grid, +) +from spatialdata_io.experimental.regular_grid import ( + DEFAULT_MAX_ROW_GROUPS_PER_FILE, + RegularGrid, +) +from spatialdata_io.experimental.shapes_parquet import ( + write_cell_metadata, + write_shapes_regular_grid, +) + +__all__ = ["add_spatial_tiling", "xenium_spatially_tiled"] + +#: Directory inside the store holding derived (non-canonical) profile assets. +PROFILE_DIR = "visualization" + +#: Display colours cycled through when a channel has none assigned. First is blue, which +#: is the conventional nuclear stain colour and usually channel 0 (DAPI). +_DEFAULT_CHANNEL_COLORS = [ + (0, 0, 255), + (0, 255, 0), + (255, 0, 0), + (255, 255, 0), + (255, 0, 255), + (0, 255, 255), +] + + +def _channels_of(element: Any) -> list[Any]: + """List an image element's channel names, falling back to indices.""" + level = element[next(iter(element.children))] if hasattr(element, "children") else element + array = level[next(iter(level.data_vars))] if hasattr(level, "data_vars") else level + coords = getattr(array, "coords", {}) + if "c" in coords: + return [str(c) for c in coords["c"].values] + size = dict(zip(array.dims, array.shape, strict=True)).get("c", 1) + return list(range(size)) + + +def _channel_label(channel: Any, index: int) -> str: + """A filesystem- and URL-safe label for a channel. + + Xenium channel names include slashes ('ATP1A1/CD45/E-Cadherin'), which would other- + wise create nested directories and break the manifest's relative paths. + """ + if isinstance(channel, int): + return f"channel_{channel}" + safe = "".join(ch if ch.isalnum() else "_" for ch in str(channel)).strip("_").lower() + while "__" in safe: + safe = safe.replace("__", "_") + return safe or f"channel_{index}" + + +def _grid_for(points: Any, transform: DisplayTransform, tile_size_px: float) -> RegularGrid: + """Derive the grid covering the points element in display pixel space.""" + x = points["x"].max().compute() if hasattr(points["x"].max(), "compute") else points["x"].max() + y = points["y"].max().compute() if hasattr(points["y"].max(), "compute") else points["y"].max() + px, py = transform.apply(np.array([float(x)]), np.array([float(y)])) + return RegularGrid.from_bounds(0, 0, float(np.rint(px[0])), float(np.rint(py[0])), tile_size_px) + + +def add_spatial_tiling( + store: str | Path, + *, + points_element: str = "transcripts", + shapes_element: str | None = "cell_boundaries", + table_element: str | None = "table", + coordinate_system: str = "global", + tile_size_px: float = 250.0, + max_row_groups_per_file: int = DEFAULT_MAX_ROW_GROUPS_PER_FILE, + feature_key: str = "feature_name", + technology: str = "Xenium", + include_cbg: bool = True, + image_element: str | None = None, + image_channels: list[int | str] | None = None, + image_colors: dict[str, tuple[int, int, int]] | None = None, + image_tile_size: int = 512, + compression: str = "zstd", +) -> dict[str, Any]: + """Add the regular-grid visualization profile to an existing SpatialData store. + + Re-runnable: running it again replaces the render columns and derived assets rather + than duplicating them. + + Parameters + ---------- + store + Path to the ``.zarr`` store to tile, modified in place. + points_element + Name of the Points element holding transcripts. + shapes_element + Name of the Shapes element holding cell boundaries, or ``None`` to skip. + table_element + Name of the annotating table, used for the gene order and the CBG. ``None`` skips + the CBG and derives feature codes from the observed features alone. + coordinate_system + Coordinate system defining display pixel space. + tile_size_px + Tile edge length in display pixels. The default of 250 gives roughly 20 cells per + tile on Xenium-density tissue, which is the granularity the viewer fetches at. + max_row_groups_per_file + Row groups per chunk file. + feature_key + Column in the points element holding the feature name. + technology + Celldega technology string recorded in the manifest. + include_cbg + Whether to write the gene-major cell-by-gene files. + image_element + Name of an image element to render into a WebP display pyramid, or ``None`` to + skip images. The canonical OME-Zarr image is left untouched either way. + image_channels + Channels to render, by name or index. Defaults to every channel in the element. + image_colors + Optional display colour per channel name. Channels without an entry get a colour + from a default palette. + image_tile_size + Image tile edge length in pixels. + compression + Parquet compression codec. + + Returns + ------- + The profile manifest. + """ + import spatialdata + + store = Path(store) + sdata = spatialdata.read_zarr(store) + + if points_element not in sdata.points: + raise ValueError(f"points element {points_element!r} not found; have {list(sdata.points)}") + points = sdata.points[points_element] + table = sdata.tables[table_element] if table_element else None + + transform = DisplayTransform.from_element(points, coordinate_system) + grid = _grid_for(points, transform, tile_size_px) + + if table is not None: + catalog = FeatureCatalog.from_points_and_table(points, table, feature_key=feature_key) + else: + observed = points[feature_key] + names = sorted(observed.cat.as_known().cat.categories) if hasattr(observed, "cat") else [] + catalog = FeatureCatalog(names=tuple(names), n_genes=len(names)) + + profile_dir = store / PROFILE_DIR / PROFILE_NAME + profile_dir.mkdir(parents=True, exist_ok=True) + + # The render columns go to a standalone file inside the profile directory. A viewer + # reads every column of it, so it needs no column projection, and the canonical + # element is left free of nested Arrow columns. + transcripts = write_points_regular_grid( + points, + profile_dir / "trx", + catalog=catalog, + grid=grid, + display_transform=transform, + feature_key=feature_key, + max_row_groups_per_file=max_row_groups_per_file, + compression=compression, + render_only=True, + overwrite=True, + ) + # The canonical element is re-ordered into tile row groups but keeps only its own + # columns, so it still round-trips through SpatialData.write() and normal reads. + write_points_regular_grid( + points, + store / "points" / points_element / "points.parquet", + catalog=catalog, + grid=grid, + display_transform=transform, + feature_key=feature_key, + max_row_groups_per_file=max_row_groups_per_file, + compression=compression, + overwrite=True, + ) + + cell_segmentation = None + cell_metadata = None + cell_names: list[str] | None = None + if shapes_element: + if shapes_element not in sdata.shapes: + raise ValueError(f"shapes element {shapes_element!r} not found; have {list(sdata.shapes)}") + shapes = sdata.shapes[shapes_element] + shapes_transform = DisplayTransform.from_element(shapes, coordinate_system) + # Cells are the overview representation, so the client needs every centroid up + # front. Centroids are not stored anywhere in the SpatialData store (the Xenium + # reader puts no x/y_centroid in obs), so they are derived from the geometry -- + # the same centroids already computed for tile assignment. + cell_metadata = write_cell_metadata( + shapes, + profile_dir / "cell_metadata.parquet", + display_transform=shapes_transform, + cell_index=list(table.obs_names) if table is not None else None, + ) + cell_names = [str(k) for k in (table.obs_names if table is not None else shapes.index)] + cell_segmentation = write_shapes_regular_grid( + shapes, + profile_dir / "cell_seg", + grid=grid, + display_transform=shapes_transform, + cell_index=list(table.obs_names) if table is not None else None, + max_row_groups_per_file=max_row_groups_per_file, + compression=compression, + render_only=True, + overwrite=True, + ) + write_shapes_regular_grid( + shapes, + store / "shapes" / shapes_element / "shapes.parquet", + grid=grid, + display_transform=shapes_transform, + cell_index=list(table.obs_names) if table is not None else None, + max_row_groups_per_file=max_row_groups_per_file, + compression=compression, + overwrite=True, + ) + + cbg = None + if include_cbg and table is not None: + cbg = write_cbg_row_groups( + table, + profile_dir / "cbg", + catalog=catalog, + max_row_groups_per_file=max_row_groups_per_file, + compression=compression, + overwrite=True, + ) + + images: dict[str, Any] = {} + image_info: list[dict[str, Any]] = [] + image_dimensions: dict[str, Any] | None = None + max_pyramid_zoom: int | None = None + if image_element: + from spatialdata_io.experimental.webp_parquet import write_webp_pyramid + + if image_element not in sdata.images: + raise ValueError(f"image element {image_element!r} not found; have {list(sdata.images)}") + element = sdata.images[image_element] + channels = image_channels if image_channels is not None else _channels_of(element) + + for index, channel in enumerate(channels): + label = _channel_label(channel, index) + pyramid = write_webp_pyramid( + element, + profile_dir / "images" / label, + channel=channel, + tile_size=image_tile_size, + source_element=image_element, + overwrite=True, + ) + # ImageRowGroupReader resolves files as baseUrl/directory/file and reads the + # per-zoom grid from the entry's zoom_info. + pyramid["directory"] = f"images/{label}" + images[label] = pyramid + colour = (image_colors or {}).get(label) or _DEFAULT_CHANNEL_COLORS[index % len(_DEFAULT_CHANNEL_COLORS)] + image_info.append({"name": label, "button_name": str(channel), "color": list(colour)}) + # Every channel of one element shares its dimensions and pyramid depth. + image_dimensions = { + "width": pyramid["source_width"], + "height": pyramid["source_height"], + "tile_size": image_tile_size, + } + max_pyramid_zoom = pyramid["max_zoom"] + + # The gene name is the index and must be preserved: a client reads the gene list from + # it, and without it the viewer simply shows no transcript controls at all. + catalog.to_frame(table.X if table is not None else None).to_parquet(profile_dir / "meta_gene.parquet") + + # Files a client reads at fixed paths rather than through the manifest. Writing them + # is what lets the profile directory stand in for a DegaFiles root, so no client needs + # to know it is looking at a SpatialData store. + _write_micron_to_image_transform(profile_dir / "micron_to_image_transform.csv", transform) + cluster_info = None + if cell_names is not None: + cluster_info = _write_cell_clusters(profile_dir / "cell_clusters", cell_names, None) + + manifest = build_manifest( + grid=grid, + technology=technology, + transcripts=transcripts, + cell_segmentation=cell_segmentation, + cbg=cbg, + images=images, + image_info=image_info, + image_dimensions=image_dimensions, + max_pyramid_zoom=max_pyramid_zoom, + fixed_path_assets={ + "meta_gene": "meta_gene.parquet", + "micron_to_image_transform": "micron_to_image_transform.csv", + **({"cell_metadata": cell_metadata} if cell_metadata else {}), + **({"cell_clusters": cluster_info} if cluster_info else {}), + }, + source={ + "store": store.name, + "points_element": points_element, + "shapes_element": shapes_element, + "table_element": table_element, + "image_element": image_element, + "coordinate_system": coordinate_system, + "tile_size_px": tile_size_px, + }, + ) + validate_manifest(manifest, base_path=profile_dir) + write_manifest(manifest, profile_dir) + return manifest + + +def xenium_spatially_tiled( + raw_path: str | Path, + output_path: str | Path, + *, + tile_size_px: float = 250.0, + max_row_groups_per_file: int = DEFAULT_MAX_ROW_GROUPS_PER_FILE, + include_cbg: bool = True, + compression: str = "zstd", + overwrite: bool = False, + 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. + include_cbg + Whether to write the gene-major cell-by-gene files. + compression + Parquet compression codec. + overwrite + Replace ``output_path`` if it exists. + tiling + Extra keyword arguments for :func:`add_spatial_tiling`, for example + ``{"image_element": "morphology_focus", "image_channel": "DAPI"}``. Kept separate + from ``xenium_kwargs`` because the two functions have distinct option sets. + xenium_kwargs + Forwarded to :func:`spatialdata_io.xenium`. + + Returns + ------- + The profile manifest. + """ + from spatialdata_io.readers.xenium import xenium + + output_path = Path(output_path) + if output_path.exists(): + if not overwrite: + raise FileExistsError(f"{output_path} exists; pass overwrite=True to replace it") + shutil.rmtree(output_path) + + sdata = xenium(raw_path, **xenium_kwargs) + sdata.write(output_path) + return add_spatial_tiling( + output_path, + tile_size_px=tile_size_px, + max_row_groups_per_file=max_row_groups_per_file, + include_cbg=include_cbg, + compression=compression, + **(tiling or {}), + ) + + +def _write_micron_to_image_transform(path: Path, transform: DisplayTransform) -> None: + """Write the micron-to-image affine Celldega reads at a fixed path. + + Coordinates in the profile are already in display pixels, so this is not needed to + place anything. It is needed for the scale bar and any physical-units readout, so it + must be the real micron-to-pixel affine and not identity -- writing identity would + make the scale bar wrong by the pixel size. + """ + (a, b, c), (d, e, f) = transform.matrix + rows = [f"{a} {b} {c}", f"{d} {e} {f}", "0.0 0.0 1.0"] + path.write_text("\n".join(rows) + "\n") + + +def _write_cell_clusters(directory: Path, cell_names: list[str], clusters: Any | None) -> dict[str, Any]: + """Write the cluster assignment and palette Celldega reads at a fixed path. + + SpatialData does not require a clustering, and the Xenium reader does not load one, so + when none is supplied every cell is placed in a single group. That keeps the viewer's + category machinery working instead of failing on a missing file; it is a placeholder, + not a scientific result. + """ + import colorsys + + import pandas as pd + + directory.mkdir(parents=True, exist_ok=True) + if clusters is None: + labels = pd.Series(["unclustered"] * len(cell_names), index=cell_names, dtype=object) + else: + labels = pd.Series([str(v) for v in clusters], index=cell_names, dtype=object) + + counts = labels.value_counts() + palette = {} + for i, name in enumerate(counts.index): + r, g, b = colorsys.hsv_to_rgb((i * 0.618033988749895) % 1.0, 0.6, 0.9) + palette[name] = f"#{int(r * 255):02x}{int(g * 255):02x}{int(b * 255):02x}" + + pd.DataFrame({"cluster": labels}).to_parquet(directory / "cluster.parquet") + pd.DataFrame( + {"color": [palette[n] for n in counts.index], "count": counts.to_numpy()}, + index=list(counts.index), + ).to_parquet(directory / "meta_cluster.parquet") + + return {"n_clusters": int(counts.size), "placeholder": clusters is None} diff --git a/src/spatialdata_io/experimental/webp_parquet.py b/src/spatialdata_io/experimental/webp_parquet.py new file mode 100644 index 00000000..dbf3aaeb --- /dev/null +++ b/src/spatialdata_io/experimental/webp_parquet.py @@ -0,0 +1,322 @@ +"""Write a browser-ready WebP image pyramid as Parquet row groups. + +The canonical OME-Zarr image stays authoritative; this is a derived *display cache*, and +must never be treated as the quantitative image. + +Layout matches Celldega's existing ``ImageRowGroupReader``: columns ``zoom``, ``tile_x``, +``tile_y``, ``image_data`` (encoded WebP bytes), one tile per row group, with a +``zoom_info`` map in the Parquet schema metadata:: + + row_group_index = zoom_info[zoom].row_group_offset + tile_x * num_tiles_y + tile_y + +Zoom numbering follows the DeepZoom convention, so the same numbers work against tiles +produced by ``vips dzsave``: level ``max_zoom`` is full resolution and each level below +halves both dimensions, down to a level that fits in a single tile. + +Encoding uses Pillow rather than libvips, so this adds no system dependency. Pyramid +levels are produced from the SpatialData image's own multiscale levels where they line up, +and by downsampling otherwise. +""" + +from __future__ import annotations + +import io +import json +import math +import shutil +from pathlib import Path +from typing import Any + +import numpy as np +import pyarrow as pa +import pyarrow.parquet as pq +from numpy.typing import NDArray + +__all__ = ["write_webp_pyramid", "DEFAULT_IMAGE_TILE_SIZE"] + +#: DeepZoom tile size used by Celldega's image pipeline. +DEFAULT_IMAGE_TILE_SIZE = 512 + +#: Image tiles are numerous and individually small, so more fit comfortably per file +#: than for transcripts. +DEFAULT_IMAGE_ROW_GROUPS_PER_FILE = 2000 + + +def _require_pillow() -> Any: + try: + from PIL import Image, features + except ImportError as exc: + raise RuntimeError("writing a WebP pyramid requires Pillow: pip install 'Pillow>=10'") from exc + if not features.check("webp"): + raise RuntimeError("this Pillow build has no WebP support; reinstall Pillow with WebP enabled") + return Image + + +def _select_channel(array: Any, channel: int | str | None) -> Any: + """Reduce one multiscale level to a 2D (y, x) plane.""" + if hasattr(array, "data_vars"): + array = array[next(iter(array.data_vars))] + dims = getattr(array, "dims", None) + if dims and "c" in dims: + if channel is None: + array = array.isel(c=0) + elif isinstance(channel, str): + array = array.sel(c=channel) + else: + array = array.isel(c=channel) + return array + + +def _as_2d(image: Any, channel: int | str | None) -> tuple[Any, Any]: + """Return ``(full_resolution_plane, coarse_plane)`` for a SpatialData image element. + + The coarse plane is the smallest available multiscale level, used only to choose the + display intensity window cheaply. For a single-scale image both are the same array. + """ + if hasattr(image, "children") and len(image.children): + levels = list(image.children) + return ( + _select_channel(image[levels[0]], channel), + _select_channel(image[levels[-1]], channel), + ) + plane = _select_channel(image, channel) + return plane, plane + + +def _display_window(sample_source: Any, display_min: float | None, display_max: float | None) -> tuple[float, float]: + """Choose the intensity window. + + The default is the *full dtype range* for integer images, which is a linear mapping + and no stretch at all. That deliberately matches Celldega's own image pipeline, which + saves raw values and leaves brightening to the viewer's intensity slider. A percentile + stretch here would be applied on top of that slider and blow out the mid-tones -- on + Xenium DAPI it took a tile from mean 1.3 to mean 37.4 with 1.2% of pixels saturated. + + Pass ``display_min``/``display_max`` to window explicitly; both are recorded in the + manifest so the choice is reproducible and invalidatable. + """ + if display_min is not None and display_max is not None: + lo, hi = float(display_min), float(display_max) + else: + sample = np.asarray(sample_source) + if sample.dtype.kind in "ui": + info = np.iinfo(sample.dtype) + default_lo, default_hi = float(info.min), float(info.max) + else: + finite = sample[np.isfinite(sample)] + default_lo = float(finite.min()) if finite.size else 0.0 + default_hi = float(finite.max()) if finite.size else 1.0 + lo = float(display_min) if display_min is not None else default_lo + hi = float(display_max) if display_max is not None else default_hi + return (lo, hi if hi > lo else lo + 1.0) + + +def _to_uint8(plane: NDArray[Any], lo: float, hi: float, gamma: float, block_rows: int = 2048) -> NDArray[np.uint8]: + """Window a plane to 8-bit in row blocks, avoiding a float copy of the whole image.""" + height = plane.shape[0] + out = np.empty(plane.shape, dtype=np.uint8) + inv_gamma = 1.0 / gamma + for start in range(0, height, block_rows): + stop = min(start + block_rows, height) + block = np.asarray(plane[start:stop], dtype=np.float32) + block -= lo + block /= hi - lo + np.clip(block, 0.0, 1.0, out=block) + if gamma != 1.0: + block **= inv_gamma + block *= 255.0 + block += 0.5 + out[start:stop] = block.astype(np.uint8) + return out + + +def _downsample_half(a: NDArray[np.uint8]) -> NDArray[np.uint8]: + """Box-filter by 2 in both axes, padding odd edges by replication.""" + h, w = a.shape + if h % 2: + a = np.vstack([a, a[-1:]]) + if w % 2: + a = np.hstack([a, a[:, -1:]]) + return a.reshape(a.shape[0] // 2, 2, a.shape[1] // 2, 2).mean(axis=(1, 3)).astype(np.uint8) + + +def write_webp_pyramid( + image: Any, + output_dir: str | Path, + *, + channel: int | str | None = None, + tile_size: int = DEFAULT_IMAGE_TILE_SIZE, + display_min: float | None = None, + display_max: float | None = None, + gamma: float = 1.0, + quality: int = 85, + lossless: bool = False, + max_row_groups_per_file: int = DEFAULT_IMAGE_ROW_GROUPS_PER_FILE, + source_element: str = "", + overwrite: bool = False, +) -> dict[str, Any]: + """Write a DeepZoom-numbered WebP pyramid as Parquet row groups. + + Parameters + ---------- + image + A SpatialData image element (``DataTree`` or ``DataArray``). + output_dir + Directory to write chunk files into. Written atomically. + channel + Channel index or name to render. Defaults to the first channel. + tile_size + Tile edge length in pixels. + display_min, display_max + Intensity window. Defaults to the 1st and 99.9th percentiles of the full-resolution + plane, which is a display choice and is recorded in the returned metadata. + gamma + Display gamma applied after windowing. + quality + WebP quality when ``lossless`` is False. + lossless + Whether to encode losslessly. + max_row_groups_per_file + Tiles per chunk file. + source_element + Name of the canonical image element, recorded for invalidation. + overwrite + Replace ``output_dir`` if it exists. + + Returns + ------- + The manifest fragment describing the written pyramid. + """ + Image = _require_pillow() + + output_dir = Path(output_dir) + if output_dir.exists() and not overwrite: + raise FileExistsError(f"{output_dir} exists; pass overwrite=True to replace it") + + plane, coarse = _as_2d(image, channel) + if plane.ndim != 2: + raise ValueError(f"expected a 2D plane after channel selection, got shape {plane.shape}") + source_dtype = str(plane.dtype) + + applied_min, applied_max = _display_window(coarse, display_min, display_max) + full = _to_uint8(plane, applied_min, applied_max, gamma) + height, width = full.shape + + # DeepZoom numbering: level max_zoom is full resolution and each level below halves + # both dimensions, down to level 0 (a single pixel). Every level is generated so the + # numbering matches `vips dzsave` output exactly. + max_zoom = max(1, math.ceil(math.log2(max(width, height)))) + levels: dict[int, NDArray[np.uint8]] = {max_zoom: full} + current = full + for zoom in range(max_zoom - 1, -1, -1): + current = _downsample_half(current) + levels[zoom] = current + + schema = pa.schema( + [ + pa.field("zoom", pa.int32()), + pa.field("tile_x", pa.int32()), + pa.field("tile_y", pa.int32()), + pa.field("image_data", pa.binary()), + ] + ) + + # Enumerate tiles in the reader's order: zoom ascending, then column-major within zoom. + ordered: list[tuple[int, int, int]] = [] + zoom_info: dict[str, dict[str, int]] = {} + for zoom in sorted(levels): + lh, lw = levels[zoom].shape + nx = max(1, math.ceil(lw / tile_size)) + ny = max(1, math.ceil(lh / tile_size)) + zoom_info[str(zoom)] = { + "num_tiles_x": nx, + "num_tiles_y": ny, + "num_tiles": nx * ny, + "row_group_offset": len(ordered), + } + ordered.extend((zoom, tx, ty) for tx in range(nx) for ty in range(ny)) + + n_files = max(1, -(-len(ordered) // max_row_groups_per_file)) + width_digits = len(str(n_files - 1)) if n_files > 1 else 1 + filenames = [f"chunk_{i:0{width_digits}d}.parquet" for i in range(n_files)] + + schema = schema.with_metadata( + { + b"zoom_info": json.dumps(zoom_info).encode(), + b"storage_mode": b"row_groups_image_chunked", + b"max_row_groups_per_file": str(max_row_groups_per_file).encode(), + b"tile_size": str(tile_size).encode(), + b"profile": b"grid_files_v1", + } + ) + + staging = output_dir.with_name(output_dir.name + ".tmp") + if staging.exists(): + shutil.rmtree(staging) + staging.mkdir(parents=True) + + encode_kwargs: dict[str, Any] = {"format": "WEBP", "lossless": lossless} + if not lossless: + encode_kwargs["quality"] = quality + + try: + writer: pq.ParquetWriter | None = None + current_file = -1 + for index, (zoom, tx, ty) in enumerate(ordered): + file_index = index // max_row_groups_per_file + if file_index != current_file: + if writer is not None: + writer.close() + writer = pq.ParquetWriter(staging / filenames[file_index], schema, write_statistics=False) + current_file = file_index + + level = levels[zoom] + crop = level[ty * tile_size : (ty + 1) * tile_size, tx * tile_size : (tx + 1) * tile_size] + buf = io.BytesIO() + Image.fromarray(crop, mode="L").save(buf, **encode_kwargs) + + assert writer is not None + writer.write_table( + pa.table( + { + "zoom": pa.array([zoom], pa.int32()), + "tile_x": pa.array([tx], pa.int32()), + "tile_y": pa.array([ty], pa.int32()), + "image_data": pa.array([buf.getvalue()], pa.binary()), + }, + schema=schema, + ) + ) + if writer is not None: + writer.close() + except BaseException: + shutil.rmtree(staging, ignore_errors=True) + raise + + if output_dir.exists(): + shutil.rmtree(output_dir) + staging.rename(output_dir) + + return { + "directory": output_dir.name, + "files": filenames, + "max_row_groups_per_file": max_row_groups_per_file, + "total_row_groups": len(ordered), + "zoom_info": zoom_info, + "min_zoom": min(levels), + "max_zoom": max_zoom, + "tile_size": tile_size, + "image_format": ".webp", + # Recorded so the display cache can be invalidated when any of it changes. + "source_element": source_element, + "source_width": int(width), + "source_height": int(height), + "source_dtype": source_dtype, + "channel": channel, + "display_min": applied_min, + "display_max": applied_max, + "gamma": gamma, + "downsampling": "box-2x2-mean", + "webp_lossless": lossless, + "webp_quality": None if lossless else quality, + } diff --git a/tests/test_cbg_parquet.py b/tests/test_cbg_parquet.py new file mode 100644 index 00000000..ee5411f1 --- /dev/null +++ b/tests/test_cbg_parquet.py @@ -0,0 +1,190 @@ +"""Tests for the gene-major cell-by-gene writer.""" + +from __future__ import annotations + +import json +from pathlib import Path + +import numpy as np +import pandas as pd +import pyarrow.parquet as pq +import pytest +import scipy.sparse as sp +from anndata import AnnData + +from spatialdata_io.experimental.cbg_parquet import write_cbg_row_groups +from spatialdata_io.experimental.feature_catalog import FeatureCatalog + +GENES = ["GENEA", "GENEB", "GENEC"] +CELLS = ["cell-0", "cell-1", "cell-2", "cell-3"] + +#: rows = cells, cols = genes. GENEC is deliberately all-zero. +DENSE = np.array( + [ + [5.0, 0.0, 0.0], + [0.0, 2.0, 0.0], + [3.0, 7.0, 0.0], + [0.0, 0.0, 0.0], + ], + dtype=np.float32, +) + + +@pytest.fixture +def table() -> AnnData: + return AnnData( + X=sp.csr_matrix(DENSE), + obs=pd.DataFrame(index=CELLS), + var=pd.DataFrame(index=GENES), + ) + + +@pytest.fixture +def catalog() -> FeatureCatalog: + return FeatureCatalog.from_features_and_table([*GENES, "NegControlProbe_0001"], GENES) + + +@pytest.fixture +def written(tmp_path: Path, table: AnnData, catalog: FeatureCatalog) -> tuple[Path, dict]: + out = tmp_path / "cbg" + manifest = write_cbg_row_groups(table, out, catalog=catalog) + return out, manifest + + +def _row_group_for(directory: Path, manifest: dict, gene: str): + rg = manifest["gene_to_row_group"][gene] + file_index, local = divmod(rg, manifest["max_row_groups_per_file"]) + return pq.ParquetFile(directory / manifest["files"][file_index]).read_row_group(local) + + +# -- layout ------------------------------------------------------------------- + + +def test_one_row_group_per_gene(written: tuple[Path, dict]) -> None: + directory, manifest = written + total = sum(pq.ParquetFile(directory / f).metadata.num_row_groups for f in manifest["files"]) + assert total == len(GENES) == manifest["num_genes"] + + +def test_row_group_index_equals_feature_code(written: tuple[Path, dict], catalog: FeatureCatalog) -> None: + """The invariant that lets one integer address both a transcript's gene and its CBG vector.""" + _, manifest = written + for gene, rg in manifest["gene_to_row_group"].items(): + assert catalog.names.index(gene) == rg + assert manifest["row_group_equals_feature_code"] is True + + +def test_controls_get_no_row_group(written: tuple[Path, dict]) -> None: + _, manifest = written + assert "NegControlProbe_0001" not in manifest["gene_to_row_group"] + assert manifest["num_genes"] == len(GENES) + + +def test_schema_matches_celldega_reader(written: tuple[Path, dict]) -> None: + directory, manifest = written + f = pq.ParquetFile(directory / manifest["files"][0]) + assert f.schema_arrow.names == ["cell_id", "expression", "gene"] + meta = f.schema_arrow.metadata + assert json.loads(meta[b"gene_to_row_group"]) == manifest["gene_to_row_group"] + assert meta[b"storage_mode"] == b"row_groups_cbg_chunked" + assert int(meta[b"num_genes"]) == len(GENES) + + +# -- values ------------------------------------------------------------------- + + +def test_sparse_values_match_the_source_matrix(written: tuple[Path, dict]) -> None: + directory, manifest = written + for col, gene in enumerate(GENES): + rg = _row_group_for(directory, manifest, gene).to_pandas() + expected = {i: DENSE[i, col] for i in range(len(CELLS)) if DENSE[i, col] != 0} + assert dict(zip(rg["cell_id"], rg["expression"], strict=True)) == expected + assert set(rg["gene"]) <= {gene} + + +def test_zero_values_are_omitted(written: tuple[Path, dict]) -> None: + directory, manifest = written + for gene in GENES: + rg = _row_group_for(directory, manifest, gene) + assert all(v != 0 for v in rg["expression"].to_pylist()) + + +def test_all_zero_gene_is_an_empty_row_group(written: tuple[Path, dict]) -> None: + """GENEC still occupies its slot so the feature_code invariant holds.""" + directory, manifest = written + assert "GENEC" in manifest["gene_to_row_group"] + assert _row_group_for(directory, manifest, "GENEC").num_rows == 0 + + +def test_explicit_stored_zeros_are_dropped(tmp_path: Path, catalog: FeatureCatalog) -> None: + """A sparse matrix can carry stored zeros; they are not expression.""" + m = sp.csr_matrix(DENSE) + m[3, 0] = 0 # creates an explicit stored zero + adata = AnnData(X=m, obs=pd.DataFrame(index=CELLS), var=pd.DataFrame(index=GENES)) + out = tmp_path / "cbg" + manifest = write_cbg_row_groups(adata, out, catalog=catalog) + rg = _row_group_for(out, manifest, "GENEA").to_pandas() + assert 3 not in set(rg["cell_id"]) + + +def test_dense_matrix_is_supported(tmp_path: Path, catalog: FeatureCatalog) -> None: + adata = AnnData(X=DENSE.copy(), obs=pd.DataFrame(index=CELLS), var=pd.DataFrame(index=GENES)) + out = tmp_path / "cbg" + manifest = write_cbg_row_groups(adata, out, catalog=catalog) + rg = _row_group_for(out, manifest, "GENEB").to_pandas() + assert dict(zip(rg["cell_id"], rg["expression"], strict=True)) == {1: 2.0, 2: 7.0} + + +def test_layer_can_be_selected(tmp_path: Path, table: AnnData, catalog: FeatureCatalog) -> None: + table.layers["scaled"] = sp.csr_matrix(DENSE * 10) + out = tmp_path / "cbg" + manifest = write_cbg_row_groups(table, out, catalog=catalog, layer="scaled") + rg = _row_group_for(out, manifest, "GENEA").to_pandas() + assert dict(zip(rg["cell_id"], rg["expression"], strict=True)) == {0: 50.0, 2: 30.0} + + +# -- cell codes --------------------------------------------------------------- + + +def test_cell_codes_default_to_table_order(written: tuple[Path, dict]) -> None: + directory, manifest = written + rg = _row_group_for(directory, manifest, "GENEA").to_pandas() + assert sorted(rg["cell_id"]) == [0, 2] + + +def test_explicit_cell_codes_are_honoured(tmp_path: Path, table: AnnData, catalog: FeatureCatalog) -> None: + codes = {name: i for i, name in enumerate(reversed(CELLS))} + out = tmp_path / "cbg" + manifest = write_cbg_row_groups(table, out, catalog=catalog, cell_codes=codes) + rg = _row_group_for(out, manifest, "GENEA").to_pandas() + assert sorted(rg["cell_id"]) == sorted([codes["cell-0"], codes["cell-2"]]) + + +def test_missing_cell_code_is_reported(tmp_path: Path, table: AnnData, catalog: FeatureCatalog) -> None: + with pytest.raises(ValueError, match="no cell_code"): + write_cbg_row_groups(table, tmp_path / "cbg", catalog=catalog, cell_codes={"cell-0": 0}) + + +def test_gene_missing_from_table_is_reported(tmp_path: Path, table: AnnData) -> None: + bad = FeatureCatalog(names=("GENEA", "GHOST"), n_genes=2) + with pytest.raises(ValueError, match="absent from table.var_names"): + write_cbg_row_groups(table, tmp_path / "cbg", catalog=bad) + + +# -- chunking ----------------------------------------------------------------- + + +def test_multi_file_chunking(tmp_path: Path, table: AnnData, catalog: FeatureCatalog) -> None: + out = tmp_path / "cbg" + manifest = write_cbg_row_groups(table, out, catalog=catalog, max_row_groups_per_file=2) + assert manifest["files"] == ["chunk_0.parquet", "chunk_1.parquet"] + assert [pq.ParquetFile(out / f).metadata.num_row_groups for f in manifest["files"]] == [2, 1] + # the mapping must still resolve through the file split + rg = _row_group_for(out, manifest, "GENEC") + assert rg.num_rows == 0 + + +def test_overwrite_guard(written: tuple[Path, dict], table: AnnData, catalog: FeatureCatalog) -> None: + directory, _ = written + with pytest.raises(FileExistsError): + write_cbg_row_groups(table, directory, catalog=catalog) diff --git a/tests/test_points_parquet.py b/tests/test_points_parquet.py new file mode 100644 index 00000000..e3137f77 --- /dev/null +++ b/tests/test_points_parquet.py @@ -0,0 +1,410 @@ +"""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, + 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_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..0d3976c9 --- /dev/null +++ b/tests/test_regular_grid.py @@ -0,0 +1,277 @@ +"""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) + + +def test_meta_gene_index_is_unnamed_so_it_serializes_as_index_level_0(tmp_path) -> None: + """A client finds the gene list only under '__index_level_0__'. + + pandas writes a *named* index as a column of that name, which the client does not + look for, leaving the gene list empty and the viewer with no transcript controls. + """ + import pyarrow.parquet as pq + + from spatialdata_io.experimental.feature_catalog import FeatureCatalog + + catalog = FeatureCatalog(names=("GENEA", "GENEB", "NegControlProbe_1"), n_genes=2) + path = tmp_path / "meta_gene.parquet" + catalog.to_frame().to_parquet(path) + + names = pq.read_table(path).schema.names + assert "__index_level_0__" in names + assert "color" in names + for stat in ("mean", "std", "max", "non-zero"): + assert stat in names + + +def test_meta_gene_row_order_is_the_feature_code(tmp_path) -> None: + """A client derives its integer gene id from row position in this file. + + It then colours transcripts by indexing an array built from that position, so the + order must match feature_code exactly. Sorting the frame breaks the correspondence + and the transcripts render transparent rather than erroring. + """ + import pyarrow.parquet as pq + + from spatialdata_io.experimental.feature_catalog import FeatureCatalog + + # var_names order is deliberately NOT alphabetical, and a control sorts in the middle. + catalog = FeatureCatalog.from_features_and_table(["ZED", "ABC", "MID_Control"], var_names=["ZED", "ABC"]) + path = tmp_path / "meta_gene.parquet" + catalog.to_frame().to_parquet(path) + + df = pq.read_table(path).to_pandas() + order = list(df["__index_level_0__"]) if "__index_level_0__" in df.columns else list(df.index) + assert order == list(catalog.names) + for position, name in enumerate(order): + assert catalog.names.index(name) == position diff --git a/tests/test_shapes_parquet.py b/tests/test_shapes_parquet.py new file mode 100644 index 00000000..4628b756 --- /dev/null +++ b/tests/test_shapes_parquet.py @@ -0,0 +1,297 @@ +"""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_cell_metadata_schema_matches_celldega(tmp_path: Path, shapes: gpd.GeoDataFrame) -> None: + from spatialdata_io.experimental.shapes_parquet import write_cell_metadata + + out = tmp_path / "cell_metadata.parquet" + info = write_cell_metadata(shapes, out, display_transform=XFORM) + t = pq.read_table(out) + assert t.schema.names == ["name", "geometry"] + assert pa.types.is_string(t.schema.field("name").type) + assert pa.types.is_list(t.schema.field("geometry").type) + assert info["n_cells"] == len(SHAPES_SPEC) + + +def test_cell_metadata_holds_display_pixel_centroids(tmp_path: Path, shapes: gpd.GeoDataFrame) -> None: + from spatialdata_io.experimental.shapes_parquet import write_cell_metadata + + out = tmp_path / "cm.parquet" + write_cell_metadata(shapes, out, display_transform=XFORM) + got = dict(zip(pq.read_table(out)["name"].to_pylist(), pq.read_table(out)["geometry"].to_pylist(), strict=True)) + for name, (geom, _) in SHAPES_SPEC.items(): + c = geom.centroid + assert got[name] == pytest.approx([c.x * 2, c.y * 2], abs=0.5) + + +def test_cell_metadata_row_order_is_the_cell_code(tmp_path: Path, shapes: gpd.GeoDataFrame) -> None: + """A client takes a cell's integer id from its position here, so order is the contract. + + It must match the order used for cell_code in the tiled shapes, or colouring a cell + from an expression vector would address the wrong cell. + """ + from spatialdata_io.experimental.shapes_parquet import write_cell_metadata + + table_order = list(SHAPES_SPEC)[::-1] + meta = tmp_path / "cm.parquet" + write_cell_metadata(shapes, meta, display_transform=XFORM, cell_index=table_order) + assert pq.read_table(meta)["name"].to_pylist() == table_order + + tiled = tmp_path / "s.parquet" + write_shapes_regular_grid(shapes, tiled, grid=GRID, display_transform=XFORM, cell_index=table_order) + back = gpd.read_parquet(tiled) + for position, name in enumerate(table_order): + assert back.loc[name, CELL_CODE_COLUMN] == position + + +def test_cell_metadata_reports_a_cell_with_no_shape(tmp_path: Path, shapes: gpd.GeoDataFrame) -> None: + from spatialdata_io.experimental.shapes_parquet import write_cell_metadata + + with pytest.raises(ValueError, match="have no shape"): + write_cell_metadata( + shapes, tmp_path / "cm.parquet", display_transform=XFORM, cell_index=[*SHAPES_SPEC, "ghost"] + ) + + +def test_canonical_shapes_have_no_render_columns(written: tuple[Path, dict]) -> None: + """The nested display column stays out of the GeoParquet so it still round-trips.""" + out, manifest = written + names = pq.read_table(out).column_names + assert GEOMETRY_COLUMN not in names + assert manifest["render_only"] is False + # cell_code is a plain uint32 and is harmless to keep alongside the canonical geometry + assert CELL_CODE_COLUMN in names + + +def test_render_shapes_hold_only_the_render_columns(rendered: tuple[Path, dict]) -> None: + """A viewer reads every column of this file, so no projection is needed.""" + out, manifest = rendered + assert pq.read_table(out).column_names == [GEOMETRY_COLUMN, CELL_CODE_COLUMN] + assert manifest["render_only"] is True diff --git a/tests/test_tiled_access.py b/tests/test_tiled_access.py new file mode 100644 index 00000000..e94823ad --- /dev/null +++ b/tests/test_tiled_access.py @@ -0,0 +1,432 @@ +"""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, + build_manifest, + validate_manifest, + write_manifest, +) +from spatialdata_io.experimental.regular_grid import RegularGrid +from spatialdata_io.experimental.tiled_access import add_spatial_tiling + +GRID = RegularGrid(origin_x=0.0, origin_y=0.0, tile_size_px=10.0, num_tiles_x=2, num_tiles_y=3) + + +# -- manifest ----------------------------------------------------------------- + + +def _transcripts_entry(**over): + entry = { + "directory": "trx", + "files": ["chunk_0.parquet"], + "max_row_groups_per_file": 400, + "total_row_groups": GRID.num_tiles, + "position_column": "display_xy", + "feature_column": "feature_code", + } + entry.update(over) + return entry + + +def test_manifest_uses_celldega_keys() -> None: + m = build_manifest(grid=GRID, transcripts=_transcripts_entry()) + for key in ("technology", "use_row_groups", "tile_grid", "row_group_files", "image_info"): + assert key in m + assert m["use_row_groups"] is True + assert m["profile"] == PROFILE_NAME + assert m["tile_grid"]["num_tiles_x"] == 2 + + +def test_manifest_omits_absent_elements() -> None: + m = build_manifest(grid=GRID, transcripts=_transcripts_entry()) + assert "cell_segmentation" not in m["row_group_files"] + assert "cbg" not in m["row_group_files"] + # images is always present so the reader's loop has something to iterate + assert m["row_group_files"]["images"] == {} + + +def test_validate_rejects_row_group_count_mismatch() -> None: + m = build_manifest(grid=GRID, transcripts=_transcripts_entry(total_row_groups=5)) + with pytest.raises(ValueError, match="does not match the tile grid"): + validate_manifest(m) + + +def test_validate_rejects_wrong_file_count() -> None: + m = build_manifest(grid=GRID, transcripts=_transcripts_entry(max_row_groups_per_file=2, files=["chunk_0.parquet"])) + with pytest.raises(ValueError, match="lists 1 file"): + validate_manifest(m) + + +def test_validate_rejects_entry_without_files_or_path() -> None: + m = build_manifest(grid=GRID, transcripts={"total_row_groups": 6, "max_row_groups_per_file": 400}) + with pytest.raises(ValueError, match="neither 'files' nor 'path'"): + validate_manifest(m) + + +def test_validate_rejects_empty_cbg_mapping() -> None: + m = build_manifest(grid=GRID, cbg={"directory": "cbg", "files": ["c.parquet"], "gene_to_row_group": {}}) + with pytest.raises(ValueError, match="gene_to_row_group is empty"): + validate_manifest(m) + + +def test_validate_rejects_cbg_row_group_out_of_range() -> None: + m = build_manifest( + grid=GRID, + cbg={ + "directory": "cbg", + "files": ["c.parquet"], + "max_row_groups_per_file": 2, + "total_row_groups": 3, + "gene_to_row_group": {"A": 0, "B": 1, "C": 9}, + }, + ) + with pytest.raises(ValueError, match="references row group 9"): + validate_manifest(m) + + +def test_validate_checks_files_exist(tmp_path: Path) -> None: + m = build_manifest(grid=GRID, transcripts=_transcripts_entry()) + with pytest.raises(ValueError, match="does not exist"): + validate_manifest(m, base_path=tmp_path) + + +def test_write_manifest_round_trips(tmp_path: Path) -> None: + m = build_manifest(grid=GRID, transcripts=_transcripts_entry()) + path = write_manifest(m, tmp_path) + assert path.name == MANIFEST_FILENAME + assert json.loads(path.read_text()) == m + + +# -- end-to-end tiling -------------------------------------------------------- + + +@pytest.fixture +def store(tmp_path: Path) -> Path: + """A minimal but complete SpatialData store: points, shapes and an annotating table.""" + from spatialdata import SpatialData + from spatialdata.models import PointsModel, ShapesModel, TableModel + from spatialdata.transformations import Scale, set_transformation + + rng = np.random.default_rng(0) + n = 200 + cells = [f"cell-{i}" for i in range(12)] + genes = ["GENEA", "GENEB", "GENEC"] + + pts = pd.DataFrame( + { + "x": rng.uniform(0, 9.9, n), + "y": rng.uniform(0, 14.9, n), + "feature_name": pd.Categorical(rng.choice([*genes, "NegControlProbe_0001"], n)), + "cell_id": rng.choice(cells, n), + } + ) + points = PointsModel.parse(pts, coordinates={"x": "x", "y": "y"}, feature_key="feature_name") + set_transformation(points, Scale([2.0, 2.0], axes=("x", "y")), "global") + + gdf = gpd.GeoDataFrame( + geometry=[ + Polygon([(x, y), (x + 0.5, y), (x + 0.5, y + 0.5), (x, y + 0.5)]) + for x, y in zip(rng.uniform(0.5, 9, 12), rng.uniform(0.5, 14, 12), strict=True) + ], + index=cells, + ) + shapes = ShapesModel.parse(gdf) + set_transformation(shapes, Scale([2.0, 2.0], axes=("x", "y")), "global") + + obs = pd.DataFrame({"region": pd.Categorical(["cell_boundaries"] * 12), "instance_id": range(12)}, index=cells) + table = TableModel.parse( + AnnData( + X=sp.csr_matrix(rng.integers(0, 5, (12, 3)).astype(np.float32)), obs=obs, var=pd.DataFrame(index=genes) + ), + region="cell_boundaries", + region_key="region", + instance_key="instance_id", + ) + + path = tmp_path / "s.zarr" + SpatialData(points={"transcripts": points}, shapes={"cell_boundaries": shapes}, tables={"table": table}).write(path) + return path + + +def test_add_spatial_tiling_produces_a_valid_profile(store: Path) -> None: + manifest = add_spatial_tiling(store, tile_size_px=10.0) + profile = store / "visualization" / PROFILE_NAME + assert (profile / MANIFEST_FILENAME).exists() + assert (profile / "meta_gene.parquet").exists() + validate_manifest(manifest, base_path=profile) + + +def test_tiled_store_still_reads_with_read_zarr(store: Path) -> None: + import spatialdata + + before = spatialdata.read_zarr(store) + n_points, n_shapes = len(before.points["transcripts"]), len(before.shapes["cell_boundaries"]) + + add_spatial_tiling(store, tile_size_px=10.0) + + after = spatialdata.read_zarr(store) + assert len(after.points["transcripts"]) == n_points + assert len(after.shapes["cell_boundaries"]) == n_shapes + # canonical columns survive untouched + assert {"x", "y", "feature_name", "cell_id"} <= set(after.points["transcripts"].columns) + # ... and the render columns are NOT here: they live in the profile's own files, so + # the canonical element keeps no nested Arrow column. + assert "display_xy" not in after.points["transcripts"].columns + assert "display_geometry" not in after.shapes["cell_boundaries"].columns + + +def test_manifest_paths_resolve_from_the_profile_directory(store: Path) -> None: + """Celldega is pointed at the profile dir, so relative paths must resolve from there.""" + manifest = add_spatial_tiling(store, tile_size_px=10.0) + profile = store / "visualization" / PROFILE_NAME + trx = manifest["row_group_files"]["transcripts"] + for f in trx["files"]: + assert (profile / trx["directory"] / f).resolve().exists() + + +def test_row_group_count_matches_grid(store: Path) -> None: + manifest = add_spatial_tiling(store, tile_size_px=10.0) + grid = RegularGrid.from_manifest_dict(manifest["tile_grid"]) + profile = store / "visualization" / PROFILE_NAME + trx = manifest["row_group_files"]["transcripts"] + total = sum(pq.ParquetFile(profile / trx["directory"] / f).metadata.num_row_groups for f in trx["files"]) + assert total == grid.num_tiles == manifest["row_group_files"]["transcripts"]["total_row_groups"] + + +def test_cbg_covers_every_gene(store: Path) -> None: + manifest = add_spatial_tiling(store, tile_size_px=10.0) + cbg = manifest["row_group_files"]["cbg"] + assert set(cbg["gene_to_row_group"]) == {"GENEA", "GENEB", "GENEC"} + assert "NegControlProbe_0001" not in cbg["gene_to_row_group"] + + +def test_tiling_is_rerunnable(store: Path) -> None: + """The two-call workflow means users will re-run this; it must not accumulate state.""" + first = add_spatial_tiling(store, tile_size_px=10.0) + second = add_spatial_tiling(store, tile_size_px=10.0) + assert first["tile_grid"] == second["tile_grid"] + assert ( + first["row_group_files"]["transcripts"]["total_row_groups"] + == (second["row_group_files"]["transcripts"]["total_row_groups"]) + ) + + import spatialdata + + after = spatialdata.read_zarr(store) + cols = list(after.points["transcripts"].columns) + assert "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") + + +def test_cbg_can_be_skipped(store: Path) -> None: + manifest = add_spatial_tiling(store, tile_size_px=10.0, include_cbg=False) + assert "cbg" not in manifest["row_group_files"] + + +# -- one-shot entry point ----------------------------------------------------- + +XENIUM_DIR = os.environ.get("SPATIALDATA_IO_XENIUM_DIR") + + +@pytest.mark.skipif(not XENIUM_DIR, reason="set SPATIALDATA_IO_XENIUM_DIR to a raw Xenium output directory") +def test_xenium_spatially_tiled_end_to_end(tmp_path: Path) -> None: + """Raw Xenium to a tiled store in one call. + + Gated on real data because there is no small raw Xenium bundle to ship: the format + needs experiment.xenium, transcripts.parquet, boundaries and a feature matrix that + all agree with each other. + """ + import spatialdata + + from spatialdata_io.experimental.tiled_access import xenium_spatially_tiled + + out = tmp_path / "tiled.zarr" + manifest = xenium_spatially_tiled( + XENIUM_DIR, + out, + tile_size_px=250.0, + include_cbg=True, + nucleus_boundaries=False, + cells_labels=False, + nucleus_labels=False, + morphology_mip=False, + morphology_focus=False, + aligned_images=False, + ) + + profile = out / "visualization" / PROFILE_NAME + validate_manifest(manifest, base_path=profile) + + grid = RegularGrid.from_manifest_dict(manifest["tile_grid"]) + trx = manifest["row_group_files"]["transcripts"] + total = sum(pq.ParquetFile(profile / trx["directory"] / f).metadata.num_row_groups for f in trx["files"]) + assert total == grid.num_tiles == trx["total_row_groups"] + + # the store is still an ordinary SpatialData store + sdata = spatialdata.read_zarr(out) + points = sdata.points["transcripts"] + assert {"x", "y", "feature_name"} <= set(points.columns) + assert {"display_xy", "feature_code"} <= set(points.columns) + assert len(points) == trx["n_rows"] + + # every fixed-path asset a client reads by convention exists + for asset in ( + "landscape_parameters.json", + "cell_metadata.parquet", + "meta_gene.parquet", + "micron_to_image_transform.csv", + "cell_clusters/cluster.parquet", + ): + assert (profile / asset).exists(), asset + + +def test_xenium_spatially_tiled_refuses_to_clobber(tmp_path: Path) -> None: + """Guard runs before any reading, so it is testable without raw data.""" + from spatialdata_io.experimental.tiled_access import xenium_spatially_tiled + + existing = tmp_path / "already.zarr" + existing.mkdir() + with pytest.raises(FileExistsError, match="pass overwrite=True"): + xenium_spatially_tiled(tmp_path / "nonexistent_raw", existing) + + +def test_tiled_store_can_still_be_rewritten_with_spatialdata_write(store: Path, tmp_path: Path) -> None: + """The reason the render columns live in their own files. + + A nested Arrow column in a Points element cannot survive dask's parquet round-trip: + SpatialData.write() either fails outright or silently returns the column as a string. + Keeping the canonical element free of one means an ordinary rewrite still works. + """ + import spatialdata + + add_spatial_tiling(store, tile_size_px=10.0) + sdata = spatialdata.read_zarr(store) + + copy = tmp_path / "rewritten.zarr" + sdata.write(copy) + + back = spatialdata.read_zarr(copy) + assert len(back.points["transcripts"]) == len(sdata.points["transcripts"]) + assert set(back.points["transcripts"].columns) == set(sdata.points["transcripts"].columns) + assert len(back.shapes["cell_boundaries"]) == len(sdata.shapes["cell_boundaries"]) + + +# -- instrument independence -------------------------------------------------- + + +@pytest.fixture +def merscope_like_store(tmp_path: Path) -> Path: + """A store shaped like a different platform: other element names, other transform. + + Nothing in the profile is Xenium-specific except defaults, so tiling has to work on + any SpatialData store regardless of which reader produced it. + """ + import geopandas as gpd + import scipy.sparse as sp + from anndata import AnnData + from shapely.geometry import Polygon + from spatialdata import SpatialData + from spatialdata.models import PointsModel, ShapesModel, TableModel + from spatialdata.transformations import Scale, set_transformation + + rng = np.random.default_rng(3) + cells = [f"c{i}" for i in range(9)] + genes = ["Gad1", "Slc17a7", "Pvalb"] + + pts = pd.DataFrame( + { + "x": rng.uniform(0, 19.9, 150), + "y": rng.uniform(0, 29.9, 150), + "gene": pd.Categorical(rng.choice([*genes, "Blank-1"], 150)), + } + ) + # A different micron-to-pixel scale from Xenium's 1/0.2125. + points = PointsModel.parse(pts, coordinates={"x": "x", "y": "y"}, feature_key="gene") + set_transformation(points, Scale([5.0, 5.0], axes=("x", "y")), "global") + + gdf = gpd.GeoDataFrame( + geometry=[ + Polygon([(x, y), (x + 0.4, y), (x + 0.4, y + 0.4), (x, y + 0.4)]) + for x, y in zip(rng.uniform(1, 18, 9), rng.uniform(1, 28, 9), strict=True) + ], + index=cells, + ) + shapes = ShapesModel.parse(gdf) + set_transformation(shapes, Scale([5.0, 5.0], axes=("x", "y")), "global") + + obs = pd.DataFrame({"region": pd.Categorical(["cell_polygons"] * 9), "instance_id": range(9)}, index=cells) + table = TableModel.parse( + AnnData( + X=sp.csr_matrix(rng.integers(0, 6, (9, 3)).astype(np.float32)), + obs=obs, + var=pd.DataFrame(index=genes), + ), + region="cell_polygons", + region_key="region", + instance_key="instance_id", + ) + + path = tmp_path / "merscope_like.zarr" + SpatialData( + points={"detected_transcripts": points}, + shapes={"cell_polygons": shapes}, + tables={"table": table}, + ).write(path) + return path + + +def test_tiling_works_on_a_non_xenium_store(merscope_like_store: Path) -> None: + """Only the defaults are Xenium; the profile itself reads generic SpatialData.""" + import spatialdata + + manifest = add_spatial_tiling( + merscope_like_store, + points_element="detected_transcripts", + shapes_element="cell_polygons", + feature_key="gene", + technology="MERSCOPE", + tile_size_px=25.0, + ) + + profile = merscope_like_store / "visualization" / PROFILE_NAME + validate_manifest(manifest, base_path=profile) + assert manifest["technology"] == "MERSCOPE" + + grid = RegularGrid.from_manifest_dict(manifest["tile_grid"]) + trx = manifest["row_group_files"]["transcripts"] + total = sum(pq.ParquetFile(profile / trx["directory"] / f).metadata.num_row_groups for f in trx["files"]) + assert total == grid.num_tiles + + # the transform is the element's own, not a hardcoded Xenium pixel size + assert manifest["row_group_files"]["transcripts"]["display_transform"]["affine_matrix"][0][0] == 5.0 + + # Blank-1 is absent from var_names, so it is a control coded above every gene + cbg = manifest["row_group_files"]["cbg"] + assert set(cbg["gene_to_row_group"]) == {"Gad1", "Slc17a7", "Pvalb"} + assert "Blank-1" not in cbg["gene_to_row_group"] + + # and the store still reads normally + sdata = spatialdata.read_zarr(merscope_like_store) + assert len(sdata.points["detected_transcripts"]) == 150 + assert len(sdata.shapes["cell_polygons"]) == 9 diff --git a/tests/test_webp_parquet.py b/tests/test_webp_parquet.py new file mode 100644 index 00000000..2fe8b78f --- /dev/null +++ b/tests/test_webp_parquet.py @@ -0,0 +1,213 @@ +"""Tests for the WebP display pyramid. + +Zoom numbering follows DeepZoom so the output is interchangeable with `vips dzsave` +tiles; the numbering tests below encode that convention explicitly. +""" + +from __future__ import annotations + +import io +import json +import math +from pathlib import Path + +import numpy as np +import pyarrow.parquet as pq +import pytest +import xarray as xr + +from spatialdata_io.experimental.webp_parquet import write_webp_pyramid + +TILE = 128 +WIDTH, HEIGHT = 500, 300 + + +def _image(width: int = WIDTH, height: int = HEIGHT, channels: int = 2) -> xr.DataArray: + y, x = np.mgrid[0:height, 0:width] + base = ((np.sin(x / 17.0) * np.cos(y / 11.0) + 1) * 1000).astype(np.uint16) + stack = np.stack([base * (i + 1) for i in range(channels)]) + return xr.DataArray(stack, dims=("c", "y", "x"), coords={"c": [f"ch{i}" for i in range(channels)]}) + + +@pytest.fixture +def written(tmp_path: Path) -> tuple[Path, dict]: + out = tmp_path / "dapi" + manifest = write_webp_pyramid(_image(), out, tile_size=TILE, source_element="morphology") + return out, manifest + + +def _tile(directory: Path, manifest: dict, zoom: int, tx: int, ty: int): + zi = manifest["zoom_info"][str(zoom)] + rg = zi["row_group_offset"] + tx * zi["num_tiles_y"] + ty + file_index, local = divmod(rg, manifest["max_row_groups_per_file"]) + return pq.ParquetFile(directory / manifest["files"][file_index]).read_row_group(local) + + +# -- zoom numbering ----------------------------------------------------------- + + +def test_max_zoom_follows_deepzoom(written: tuple[Path, dict]) -> None: + _, manifest = written + assert manifest["max_zoom"] == math.ceil(math.log2(max(WIDTH, HEIGHT))) + + +def test_every_level_down_to_zero_exists(written: tuple[Path, dict]) -> None: + """vips dzsave emits levels 0..max; matching it keeps the two interchangeable.""" + _, manifest = written + assert sorted(int(z) for z in manifest["zoom_info"]) == list(range(manifest["max_zoom"] + 1)) + + +def test_level_dimensions_halve(written: tuple[Path, dict]) -> None: + _, manifest = written + max_zoom = manifest["max_zoom"] + for zoom in range(max_zoom, 0, -1): + w = math.ceil(WIDTH / 2 ** (max_zoom - zoom)) + h = math.ceil(HEIGHT / 2 ** (max_zoom - zoom)) + zi = manifest["zoom_info"][str(zoom)] + assert zi["num_tiles_x"] == max(1, math.ceil(w / TILE)), f"zoom {zoom}" + assert zi["num_tiles_y"] == max(1, math.ceil(h / TILE)), f"zoom {zoom}" + + +def test_row_group_offsets_are_cumulative(written: tuple[Path, dict]) -> None: + _, manifest = written + running = 0 + for zoom in sorted(manifest["zoom_info"], key=int): + zi = manifest["zoom_info"][zoom] + assert zi["row_group_offset"] == running + assert zi["num_tiles"] == zi["num_tiles_x"] * zi["num_tiles_y"] + running += zi["num_tiles"] + assert running == manifest["total_row_groups"] + + +def test_index_formula_addresses_the_right_tile(written: tuple[Path, dict]) -> None: + """Mirror of ImageRowGroupReader.computeRowGroupIndex.""" + directory, manifest = written + for zoom_s, zi in manifest["zoom_info"].items(): + for tx in range(zi["num_tiles_x"]): + for ty in range(zi["num_tiles_y"]): + t = _tile(directory, manifest, int(zoom_s), tx, ty) + assert (t["zoom"][0].as_py(), t["tile_x"][0].as_py(), t["tile_y"][0].as_py()) == ( + int(zoom_s), + tx, + ty, + ) + + +# -- payload ------------------------------------------------------------------ + + +def test_tiles_decode_as_webp(written: tuple[Path, dict]) -> None: + from PIL import Image + + directory, manifest = written + t = _tile(directory, manifest, manifest["max_zoom"], 0, 0) + img = Image.open(io.BytesIO(t["image_data"][0].as_py())) + assert img.format == "WEBP" + assert img.size == (TILE, TILE) + + +def test_edge_tile_is_cropped_not_padded(written: tuple[Path, dict]) -> None: + from PIL import Image + + directory, manifest = written + max_zoom = manifest["max_zoom"] + zi = manifest["zoom_info"][str(max_zoom)] + t = _tile(directory, manifest, max_zoom, zi["num_tiles_x"] - 1, 0) + img = Image.open(io.BytesIO(t["image_data"][0].as_py())) + assert img.size[0] == WIDTH - (zi["num_tiles_x"] - 1) * TILE + + +def test_schema_matches_celldega_reader(written: tuple[Path, dict]) -> None: + directory, manifest = written + schema = pq.ParquetFile(directory / manifest["files"][0]).schema_arrow + assert schema.names == ["zoom", "tile_x", "tile_y", "image_data"] + meta = schema.metadata + assert json.loads(meta[b"zoom_info"]) == manifest["zoom_info"] + assert meta[b"storage_mode"] == b"row_groups_image_chunked" + + +def test_one_tile_per_row_group(written: tuple[Path, dict]) -> None: + directory, manifest = written + for f in manifest["files"]: + md = pq.ParquetFile(directory / f).metadata + assert all(md.row_group(i).num_rows == 1 for i in range(md.num_row_groups)) + + +# -- channels and windowing --------------------------------------------------- + + +def test_channel_can_be_selected_by_name(tmp_path: Path) -> None: + a_dir = tmp_path / "a" + b_dir = tmp_path / "b" + a = write_webp_pyramid(_image(), a_dir, tile_size=TILE, channel="ch0") + b = write_webp_pyramid(_image(), b_dir, tile_size=TILE, channel="ch1") + assert a["channel"] == "ch0" and b["channel"] == "ch1" + # ch1 is twice ch0, so the encoded pixels must differ even though the window matches. + pa_bytes = _tile(a_dir, a, a["max_zoom"], 0, 0)["image_data"][0].as_py() + pb_bytes = _tile(b_dir, b, b["max_zoom"], 0, 0)["image_data"][0].as_py() + assert pa_bytes != pb_bytes + + +def test_default_window_is_the_full_dtype_range(tmp_path: Path) -> None: + """No stretch by default, matching Celldega's pipeline. + + The viewer applies its own intensity slider; a percentile stretch here would be + applied on top of it and blow out the mid-tones. + """ + m = write_webp_pyramid(_image(), tmp_path / "d", tile_size=TILE) + assert (m["display_min"], m["display_max"]) == (0.0, float(np.iinfo(np.uint16).max)) + + +def test_explicit_window_is_recorded_and_used(tmp_path: Path) -> None: + m = write_webp_pyramid(_image(), tmp_path / "w", tile_size=TILE, display_min=100, display_max=900) + assert (m["display_min"], m["display_max"]) == (100.0, 900.0) + + +def test_degenerate_window_does_not_divide_by_zero(tmp_path: Path) -> None: + flat = xr.DataArray(np.full((1, 200, 200), 7, dtype=np.uint16), dims=("c", "y", "x")) + m = write_webp_pyramid(flat, tmp_path / "flat", tile_size=TILE) + assert m["display_max"] > m["display_min"] + + +def test_multiscale_input_uses_full_resolution(tmp_path: Path) -> None: + """A DataTree must be tiled from scale0, not from a downsampled level.""" + from xarray import DataTree + + full = _image(400, 200, channels=1) + half = full.isel(y=slice(None, None, 2), x=slice(None, None, 2)) + tree = DataTree.from_dict({"scale0": full.to_dataset(name="image"), "scale1": half.to_dataset(name="image")}) + m = write_webp_pyramid(tree, tmp_path / "ms", tile_size=TILE) + assert (m["source_width"], m["source_height"]) == (400, 200) + + +# -- metadata and safety ------------------------------------------------------ + + +def test_source_metadata_is_recorded_for_invalidation(written: tuple[Path, dict]) -> None: + _, m = written + assert m["source_element"] == "morphology" + assert (m["source_width"], m["source_height"]) == (WIDTH, HEIGHT) + assert m["source_dtype"] == "uint16" + assert m["downsampling"] == "box-2x2-mean" + assert m["image_format"] == ".webp" + + +def test_lossless_is_recorded(tmp_path: Path) -> None: + m = write_webp_pyramid(_image(), tmp_path / "ll", tile_size=TILE, lossless=True) + assert m["webp_lossless"] is True + assert m["webp_quality"] is None + + +def test_overwrite_guard(written: tuple[Path, dict]) -> None: + directory, _ = written + with pytest.raises(FileExistsError): + write_webp_pyramid(_image(), directory, tile_size=TILE) + + +def test_failure_leaves_no_partial_output(tmp_path: Path) -> None: + out = tmp_path / "boom" + bad = xr.DataArray(np.zeros((2, 2, 10, 10), dtype=np.uint16), dims=("t", "c", "y", "x")) + with pytest.raises(ValueError, match="expected a 2D plane"): + write_webp_pyramid(bad, out, tile_size=TILE) + assert not out.exists() + assert not out.with_name(out.name + ".tmp").exists()