From 0783d9851f3ae485a03c31a58a303a19ad9b0291 Mon Sep 17 00:00:00 2001 From: Nicolas Fernandez Date: Sun, 26 Jul 2026 11:32:54 -0400 Subject: [PATCH] supprt zarr for xenium trx pre-processing --- src/celldega/pre/__init__.py | 14 + src/celldega/pre/run_pre_processing.py | 88 ++-- src/celldega/pre/trx_tile.py | 160 +++---- src/celldega/pre/trx_zarr.py | 393 ++++++++++++++++++ .../test_xenium_zarr_integration.py | 89 ++++ tests/unit/test_pre/test_trx_zarr.py | 346 +++++++++++++++ 6 files changed, 992 insertions(+), 98 deletions(-) create mode 100644 src/celldega/pre/trx_zarr.py create mode 100644 tests/integration/test_xenium_zarr_integration.py create mode 100644 tests/unit/test_pre/test_trx_zarr.py diff --git a/src/celldega/pre/__init__.py b/src/celldega/pre/__init__.py index 5f76671f3..91977abdb 100644 --- a/src/celldega/pre/__init__.py +++ b/src/celldega/pre/__init__.py @@ -39,6 +39,13 @@ ) from .sbg_tile import write_pseudotranscripts_from_sbg from .trx_tile import make_trx_tiles, make_trx_tiles_row_groups +from .trx_zarr import ( + is_zarr_transcript_path, + iter_zarr_transcript_batches, + load_zarr_transcripts, + make_trx_tiles_from_zarr, + read_zarr_transcript_metadata, +) def main(*args, **kwargs): @@ -1812,11 +1819,18 @@ def add_clustering_from_adata( "add_clustering_from_adata", "boundary_tile", "get_image_info", + "is_zarr_transcript_path", + "iter_zarr_transcript_batches", "landscape", + "load_zarr_transcripts", "main", "make_trx_tiles", + "make_trx_tiles_from_zarr", + "make_trx_tiles_row_groups", "read_cbg_mtx", + "read_zarr_transcript_metadata", "resolve_xenium_morphology_ome_path", "trx_tile", + "trx_zarr", "write_identity_transform", ] diff --git a/src/celldega/pre/run_pre_processing.py b/src/celldega/pre/run_pre_processing.py index f44ddd24d..6aa77f002 100644 --- a/src/celldega/pre/run_pre_processing.py +++ b/src/celldega/pre/run_pre_processing.py @@ -95,6 +95,7 @@ def _setup_preprocessing_paths(technology, path_dega_files, data_dir, sample=Non "meta_cell_image": landscape_path / "cell_metadata.parquet", "meta_gene": landscape_path / "meta_gene.parquet", "transcripts": data_path / "transcripts.parquet", + "transcripts_zarr": data_path / "transcripts.zarr.zip", "transcript_tiles": landscape_path / "transcript_tiles", "cell_boundaries": data_path / "cell_boundaries.parquet", "cell_segmentation": landscape_path / "cell_segmentation", @@ -361,22 +362,39 @@ def make_column_names_unique_fast(df): need_trx_tiles = not trx_output_dir.exists() need_boundaries = not cell_output_dir.exists() + use_zarr_trx = technology == "Xenium" and Path(paths["transcripts_zarr"]).exists() + if need_trx_tiles: - print("\n======== Transcript Tiles (Row Groups) ========") - tile_bounds, tile_grid_info, trx_chunk_info = dega.pre.make_trx_tiles_row_groups( - technology, - str(paths["transcripts"]), - str(transform_out), - str(trx_output_dir), - coarse_tile_factor=10, - tile_size=tile_size, - chunk_size=100000, - verbose=False, - image_scale=1, - max_workers=max_workers, - path_dega_files=path_dega_files, - max_row_groups_per_file=max_row_groups_per_file, - ) + if use_zarr_trx: + print("\n======== Transcript Tiles (Row Groups, Zarr grid) ========") + tile_bounds, tile_grid_info, trx_chunk_info = dega.pre.make_trx_tiles_from_zarr( + str(paths["transcripts_zarr"]), + str(transform_out), + str(trx_output_dir), + tile_size=tile_size, + image_scale=1, + use_row_groups=True, + max_row_groups_per_file=max_row_groups_per_file, + path_dega_files=path_dega_files, + ) + else: + print("\n======== Transcript Tiles (Row Groups) ========") + tile_bounds, tile_grid_info, trx_chunk_info = ( + dega.pre.make_trx_tiles_row_groups( + technology, + str(paths["transcripts"]), + str(transform_out), + str(trx_output_dir), + coarse_tile_factor=10, + tile_size=tile_size, + chunk_size=100000, + verbose=False, + image_scale=1, + max_workers=max_workers, + path_dega_files=path_dega_files, + max_row_groups_per_file=max_row_groups_per_file, + ) + ) print(f"tile bounds: {tile_bounds}") else: print("Skipping transcript tiles, output already exists") @@ -403,20 +421,34 @@ def make_column_names_unique_fast(df): need_trx_tiles = not _output_exists(paths["transcript_tiles"]) need_boundaries = not _output_exists(paths["cell_segmentation"]) + use_zarr_trx = technology == "Xenium" and Path(paths["transcripts_zarr"]).exists() + if need_trx_tiles or need_boundaries: - print("\n======== Transcript Tiles========") - tile_bounds = dega.pre.make_trx_tiles( - technology, - str(paths["transcripts"]), - str(transform_out), - str(paths["transcript_tiles"]), - coarse_tile_factor=10, - tile_size=tile_size, - chunk_size=100000, - verbose=False, - image_scale=1, - max_workers=max_workers, - ) + if use_zarr_trx: + print("\n======== Transcript Tiles (Zarr grid) ========") + tile_bounds, _, _ = dega.pre.make_trx_tiles_from_zarr( + str(paths["transcripts_zarr"]), + str(transform_out), + str(paths["transcript_tiles"]), + tile_size=tile_size, + image_scale=1, + use_row_groups=False, + path_dega_files=path_dega_files, + ) + else: + print("\n======== Transcript Tiles========") + tile_bounds = dega.pre.make_trx_tiles( + technology, + str(paths["transcripts"]), + str(transform_out), + str(paths["transcript_tiles"]), + coarse_tile_factor=10, + tile_size=tile_size, + chunk_size=100000, + verbose=False, + image_scale=1, + max_workers=max_workers, + ) print(f"tile bounds: {tile_bounds}") else: print("Skipping transcript tiles, output already exists") diff --git a/src/celldega/pre/trx_tile.py b/src/celldega/pre/trx_tile.py index c0ffca269..a67b92815 100644 --- a/src/celldega/pre/trx_tile.py +++ b/src/celldega/pre/trx_tile.py @@ -19,6 +19,84 @@ # to avoid OOM from ``partition_by`` materializing every non-empty tile at once. STREAMING_TILE_ASSIGN_ROW_THRESHOLD = 500_000 +# Columns that only exist to build ``geometry`` / carry source identifiers and are not +# written to the tile parquets. +_TILE_DROP_COLUMNS = ("transformed_x", "transformed_y", "cell_id", "transcript_id") + + +def _transform_batch(chunk, sparse_matrix, image_scale): + """Apply the affine transform to a single polars batch. + + The input ``chunk`` must contain ``x`` and ``y`` columns (micron space). The + returned frame replaces them with rounded ``transformed_x``/``transformed_y`` + columns (image space) and keeps every other column untouched. + + Parameters + ---------- + chunk : polars.DataFrame + Batch of transcripts with ``x`` and ``y`` columns. + sparse_matrix : scipy.sparse matrix + Affine transformation matrix (micron -> image space). + image_scale : float + Scale factor applied after the affine transform. + + Returns + ------- + polars.DataFrame + ``chunk`` with ``x``/``y`` replaced by ``transformed_x``/``transformed_y``. + """ + points = np.hstack([chunk.select(["x", "y"]).to_numpy(), np.ones((chunk.height, 1))]) + transformed_points = sparse_matrix.dot(points.T).T[:, :2] + + return chunk.with_columns( + [ + (pl.Series(transformed_points[:, 0]) * image_scale).round(2).alias("transformed_x"), + (pl.Series(transformed_points[:, 1]) * image_scale).round(2).alias("transformed_y"), + ] + ).drop(["x", "y"]) + + +def _assign_tile_indices_and_geometry(trx, x_min, y_min, tile_size, n_tiles_x, n_tiles_y): + """Add clamped ``tile_x``/``tile_y`` indices and a ``geometry`` list column. + + Tile indices are computed in a single vectorized pass (``O(n)``) and clamped to + the valid grid range so edge coordinates land in the last tile. The redundant + coordinate and identifier columns in :data:`_TILE_DROP_COLUMNS` are dropped. + + Parameters + ---------- + trx : polars.DataFrame + Transcripts with ``transformed_x`` and ``transformed_y`` columns. + x_min, y_min : float + Origin of the tile grid in image space. + tile_size : float + Size of each fine tile. + n_tiles_x, n_tiles_y : int + Number of tiles along each axis (used to clamp edge coordinates). + + Returns + ------- + polars.DataFrame + Frame with ``tile_x``, ``tile_y`` and ``geometry`` columns. + """ + trx = trx.with_columns( + [ + ((pl.col("transformed_x") - x_min) / tile_size).floor().cast(pl.Int32).alias("tile_x"), + ((pl.col("transformed_y") - y_min) / tile_size).floor().cast(pl.Int32).alias("tile_y"), + ] + ) + trx = trx.with_columns( + [ + pl.col("tile_x").clip(0, n_tiles_x - 1).alias("tile_x"), + pl.col("tile_y").clip(0, n_tiles_y - 1).alias("tile_y"), + ] + ) + trx = trx.with_columns( + pl.concat_list([pl.col("transformed_x"), pl.col("transformed_y")]).alias("geometry") + ) + columns_to_drop = [col for col in _TILE_DROP_COLUMNS if col in trx.columns] + return trx.drop(columns_to_drop) + def _process_coarse_tile_transcripts( trx, @@ -244,6 +322,13 @@ def _load_transcript_data_by_technology(technology, path_trx): ).select(["name", "x", "y"]) if technology == "Xenium": + # Accept the native ``transcripts.zarr.zip`` spatial-grid bundle as an + # alternative to the flat ``transcripts.parquet``. + from .trx_zarr import is_zarr_transcript_path, load_zarr_transcripts + + if is_zarr_transcript_path(path_trx): + return load_zarr_transcripts(path_trx) + return pl.read_parquet(path_trx).select( [ pl.col("cell_id"), @@ -293,18 +378,7 @@ def _transform_coordinates_in_chunks(trx_ini, chunk_size, transformation_matrix, for start_row in tqdm(range(0, trx_ini.height, chunk_size), desc="Processing chunks"): chunk = trx_ini.slice(start_row, chunk_size) - - points = np.hstack([chunk.select(["x", "y"]).to_numpy(), np.ones((chunk.height, 1))]) - transformed_points = sparse_matrix.dot(points.T).T[:, :2] - - # Create new transformed columns and drop original x, y columns - transformed_chunk = chunk.with_columns( - [ - (pl.Series(transformed_points[:, 0]) * image_scale).round(2).alias("transformed_x"), - (pl.Series(transformed_points[:, 1]) * image_scale).round(2).alias("transformed_y"), - ] - ).drop(["x", "y"]) - all_chunks.append(transformed_chunk) + all_chunks.append(_transform_batch(chunk, sparse_matrix, image_scale)) # Concatenate all chunks after processing return pl.concat(all_chunks) @@ -328,15 +402,7 @@ def _transform_coordinates_to_parquet_shards( tqdm(range(0, trx_ini.height, chunk_size), desc="Processing chunks") ): chunk = trx_ini.slice(start_row, chunk_size) - points = np.hstack([chunk.select(["x", "y"]).to_numpy(), np.ones((chunk.height, 1))]) - transformed_points = sparse_matrix.dot(points.T).T[:, :2] - - transformed_chunk = chunk.with_columns( - [ - (pl.Series(transformed_points[:, 0]) * image_scale).round(2).alias("transformed_x"), - (pl.Series(transformed_points[:, 1]) * image_scale).round(2).alias("transformed_y"), - ] - ).drop(["x", "y"]) + transformed_chunk = _transform_batch(chunk, sparse_matrix, image_scale) row = transformed_chunk.select( [ @@ -366,27 +432,7 @@ def _spill_one_transform_shard_to_tiles( spill_root = Path(spill_root) trx = pl.read_parquet(shard_path) - trx = trx.with_columns( - [ - ((pl.col("transformed_x") - x_min) / tile_size).floor().cast(pl.Int32).alias("tile_x"), - ((pl.col("transformed_y") - y_min) / tile_size).floor().cast(pl.Int32).alias("tile_y"), - ] - ) - trx = trx.with_columns( - [ - pl.col("tile_x").clip(0, n_tiles_x - 1).alias("tile_x"), - pl.col("tile_y").clip(0, n_tiles_y - 1).alias("tile_y"), - ] - ) - trx = trx.with_columns( - pl.concat_list([pl.col("transformed_x"), pl.col("transformed_y")]).alias("geometry") - ) - columns_to_drop = [ - col - for col in ["transformed_x", "transformed_y", "cell_id", "transcript_id"] - if col in trx.columns - ] - trx = trx.drop(columns_to_drop) + trx = _assign_tile_indices_and_geometry(trx, x_min, y_min, tile_size, n_tiles_x, n_tiles_y) grouped = trx.partition_by(["tile_x", "tile_y"], as_dict=True) for (tx, ty), tile_df in grouped.items(): @@ -874,34 +920,8 @@ def _collect_tile_data_for_row_groups( print(f"Grid: {n_tiles_x} x {n_tiles_y} = {n_tiles_x * n_tiles_y} tiles") print("Calculating tile indices for all transcripts...") - # OPTIMIZED: Calculate tile indices for ALL transcripts at once (O(n)) - trx = trx.with_columns( - [ - ((pl.col("transformed_x") - x_min) / tile_size).floor().cast(pl.Int32).alias("tile_x"), - ((pl.col("transformed_y") - y_min) / tile_size).floor().cast(pl.Int32).alias("tile_y"), - ] - ) - - # Clamp to valid range (edge cases) - trx = trx.with_columns( - [ - pl.col("tile_x").clip(0, n_tiles_x - 1).alias("tile_x"), - pl.col("tile_y").clip(0, n_tiles_y - 1).alias("tile_y"), - ] - ) - - # Add geometry column - trx = trx.with_columns( - pl.concat_list([pl.col("transformed_x"), pl.col("transformed_y")]).alias("geometry") - ) - - # Drop original coordinate columns - columns_to_drop = [ - col - for col in ["transformed_x", "transformed_y", "cell_id", "transcript_id"] - if col in trx.columns - ] - trx = trx.drop(columns_to_drop) + # Calculate tile indices, clamp edges, and build geometry for ALL transcripts at once (O(n)) + trx = _assign_tile_indices_and_geometry(trx, x_min, y_min, tile_size, n_tiles_x, n_tiles_y) print("Grouping transcripts by tile...") diff --git a/src/celldega/pre/trx_zarr.py b/src/celldega/pre/trx_zarr.py new file mode 100644 index 000000000..91563bf1d --- /dev/null +++ b/src/celldega/pre/trx_zarr.py @@ -0,0 +1,393 @@ +"""Read transcripts from the 10x ``transcripts.zarr.zip`` spatial grid format. + +Modern 10x instruments (Xenium Onboard Analysis, and the Atera whole-transcriptome +preview) ship transcripts as ``transcripts.zarr.zip`` in addition to the flat +``transcripts.parquet``. The Zarr bundle stores transcripts in a ``/grids`` pyramid: +level ``0`` holds every transcript, split into spatial grid positions named +``","`` (e.g. ``grids/0/0,0``), with coarser, subsampled levels above it. + +Because the grid already partitions transcripts spatially, we can stream one grid +position at a time and feed it straight into the tiling pipeline. This avoids reading +the entire flat parquet into memory and avoids a global ``partition_by`` over every +transcript, which is the main cost of :func:`celldega.pre.trx_tile.make_trx_tiles` +on large panels. + +Reference: https://www.10xgenomics.com/support/software/xenium-onboard-analysis/latest/advanced/xoa-output-zarr +""" + +from __future__ import annotations + +from collections.abc import Iterator, Sequence +from contextlib import contextmanager +from pathlib import Path +import shutil + +import numpy as np +import polars as pl +from scipy.sparse import csr_matrix +import zarr + +from .boundary_tile import _get_name_mapping +from .trx_tile import ( + _apply_gene_mapping, + _spill_one_transform_shard_to_tiles, + _transform_batch, + _write_tiles_as_row_groups_streaming, + _write_traditional_transcript_tiles_from_spill, +) + + +# ``gene_identity`` is uint16; this sentinel marks a "no-call" (absence of a codeword). +NO_CALL_GENE_INDEX = 65535 + + +@contextmanager +def open_zarr(path: str | Path): + """Open a ``.zarr`` directory or ``.zarr.zip`` archive as a Zarr group. + + The underlying store is closed on exit, which matters for ``ZipStore`` archives + that hold an open file handle. + + Parameters + ---------- + path : str or Path + Path to a ``.zarr`` directory or ``.zarr.zip`` file. + + Yields + ------ + zarr.hierarchy.Group + The root Zarr group. + """ + path = str(path) + store = zarr.ZipStore(path, mode="r") if path.endswith(".zip") else zarr.DirectoryStore(path) + try: + yield zarr.group(store=store) + finally: + close = getattr(store, "close", None) + if close is not None: + close() + + +def is_zarr_transcript_path(path: str | Path) -> bool: + """Return ``True`` if ``path`` points at a Zarr transcript bundle.""" + return str(path).endswith((".zarr", ".zarr.zip")) + + +def _level_position_keys(grids_group, level: int, level_group) -> list[str]: + """Return the ``"col,row"`` grid-position keys for one pyramid level. + + Prefer the ``grid_keys`` attribute (authoritative ordering) and fall back to + listing the level group's subgroups for bundles that omit the attribute. + """ + grid_keys = grids_group.attrs.get("grid_keys") + if grid_keys is not None and level < len(grid_keys): + keys = [str(k) for k in grid_keys[level]] + if keys: + return keys + return sorted(level_group.group_keys()) + + +def read_zarr_transcript_metadata(path: str | Path, level: int = 0) -> dict: + """Read summary metadata from a ``transcripts.zarr.zip`` bundle. + + Parameters + ---------- + path : str or Path + Path to the Zarr transcript bundle. + level : int, optional + Pyramid level to inspect for grid positions (default ``0``, full resolution). + + Returns + ------- + dict + Keys: ``gene_names`` (list[str]), ``number_rnas`` (int or None), + ``number_levels`` (int), ``grid_size`` (list), ``position_keys`` (list[str]) + for the requested level, and ``spatial_units`` (str or None). + """ + with open_zarr(path) as root: + root_attrs = dict(root.attrs) + grids = root["grids"] + grid_attrs = dict(grids.attrs) + level_group = grids[str(level)] + position_keys = _level_position_keys(grids, level, level_group) + + return { + "gene_names": list(root_attrs.get("gene_names", [])), + "number_rnas": root_attrs.get("number_rnas"), + "number_levels": grid_attrs.get("number_levels", 1), + "grid_size": grid_attrs.get("grid_size"), + "position_keys": position_keys, + "spatial_units": root_attrs.get("spatial_units"), + } + + +def _first_column(array: np.ndarray, n_rows: int) -> np.ndarray: + """Return the first column of a possibly 2D per-transcript array.""" + values = np.asarray(array) + if values.ndim == 1: + return values + return values.reshape(n_rows, -1)[:, 0] + + +def iter_zarr_transcript_batches( + path: str | Path, + *, + level: int = 0, + drop_no_call: bool = True, + min_quality: float | None = None, + gene_names: Sequence[str] | None = None, +) -> Iterator[pl.DataFrame]: + """Yield one ``name``/``x``/``y`` batch per spatial grid position. + + Each yielded frame corresponds to a single ``grids//,`` position + and therefore covers a contiguous region of space, which keeps downstream tile + assignment local and memory bounded. + + Parameters + ---------- + path : str or Path + Path to the Zarr transcript bundle. + level : int, optional + Pyramid level to read (default ``0``, every transcript). + drop_no_call : bool, optional + Drop transcripts whose ``gene_identity`` is the no-call sentinel + (:data:`NO_CALL_GENE_INDEX`). Defaults to ``True``. + min_quality : float, optional + If given, drop transcripts with ``quality_score`` below this threshold. + gene_names : sequence of str, optional + Override the gene-name lookup table. Defaults to the bundle's ``gene_names`` + root attribute. + + Yields + ------ + polars.DataFrame + Columns ``name`` (gene name, str), ``x`` and ``y`` (micron space, float). + Empty positions are skipped. + """ + with open_zarr(path) as root: + lookup = np.asarray( + list(root.attrs.get("gene_names", [])) if gene_names is None else list(gene_names) + ) + grids = root["grids"] + level_group = grids[str(level)] + + for key in _level_position_keys(grids, level, level_group): + position = level_group[key] + location = np.asarray(position["location"][:]) + n_rows = location.shape[0] + if n_rows == 0: + continue + + gene_idx = _first_column(position["gene_identity"][:], n_rows).astype(np.int64) + + mask = np.ones(n_rows, dtype=bool) + if drop_no_call: + mask &= gene_idx != NO_CALL_GENE_INDEX + if min_quality is not None and "quality_score" in position: + quality = _first_column(position["quality_score"][:], n_rows) + mask &= quality >= min_quality + # Guard against out-of-range gene indices (e.g. control codewords). + mask &= gene_idx < lookup.shape[0] + + if not mask.any(): + continue + + yield pl.DataFrame( + { + "name": lookup[gene_idx[mask]], + "x": location[mask, 0].astype(np.float64), + "y": location[mask, 1].astype(np.float64), + } + ) + + +def load_zarr_transcripts( + path: str | Path, + *, + level: int = 0, + drop_no_call: bool = True, + min_quality: float | None = None, + gene_names: Sequence[str] | None = None, +) -> pl.DataFrame: + """Load all transcripts from a Zarr bundle into a single ``name``/``x``/``y`` frame. + + This is a convenience wrapper over :func:`iter_zarr_transcript_batches` that mirrors + the output schema of + :func:`celldega.pre.trx_tile._load_transcript_data_by_technology` for Xenium, so it + can be used as a drop-in transcript source. For large datasets prefer the streaming + tiling entry points, which never materialize the full frame. + + Parameters + ---------- + path : str or Path + Path to the Zarr transcript bundle. + level, drop_no_call, min_quality, gene_names + See :func:`iter_zarr_transcript_batches`. + + Returns + ------- + polars.DataFrame + Columns ``name``, ``x``, ``y``. Empty if the bundle has no transcripts. + """ + batches = list( + iter_zarr_transcript_batches( + path, + level=level, + drop_no_call=drop_no_call, + min_quality=min_quality, + gene_names=gene_names, + ) + ) + if not batches: + return pl.DataFrame( + {"name": [], "x": [], "y": []}, + schema={"name": pl.Utf8, "x": pl.Float64, "y": pl.Float64}, + ) + return pl.concat(batches) + + +def _resolve_gene_mapping(gene_str_to_int_mapping, path_dega_files): + """Return an explicit gene name -> int mapping, loading it if not supplied.""" + if gene_str_to_int_mapping is not None: + return gene_str_to_int_mapping + if path_dega_files is not None: + return _get_name_mapping(str(path_dega_files), layer="transcript") + return {} + + +def make_trx_tiles_from_zarr( + path_trx_zarr: str | Path, + path_transformation_matrix: str | Path, + path_output_dir: str | Path, + *, + tile_size: float = 250, + image_scale: float = 1, + use_row_groups: bool = False, + max_row_groups_per_file: int = 400, + path_dega_files: str | Path | None = None, + gene_str_to_int_mapping=None, + level: int = 0, + drop_no_call: bool = True, + min_quality: float | None = None, +): + """Tile transcripts directly from a ``transcripts.zarr.zip`` spatial grid. + + Streams one grid position at a time: apply the gene-name mapping, affine-transform + the coordinates, and spill per-tile parquet parts to disk. Because grid positions + are spatially local and processed one at a time, peak memory stays bounded and the + global ``partition_by`` used by the parquet path is avoided. + + The output layout matches the parquet-based entry points: + + * ``use_row_groups=False`` -> one ``transcripts_tile_{i}_{j}.parquet`` per non-empty + tile (same as :func:`celldega.pre.trx_tile.make_trx_tiles`). + * ``use_row_groups=True`` -> chunked row-group parquet files (same as + :func:`celldega.pre.trx_tile.make_trx_tiles_row_groups`). + + Parameters + ---------- + path_trx_zarr : str or Path + Path to the ``transcripts.zarr.zip`` bundle. + path_transformation_matrix : str or Path + Path to the whitespace-delimited micron -> image affine matrix. + path_output_dir : str or Path + Directory that will hold the tile parquet files. + tile_size : float, optional + Fine tile size in image space (default 250). + image_scale : float, optional + Scale factor applied after the affine transform (default 1). + use_row_groups : bool, optional + Write chunked row-group parquet files instead of individual tile files. + max_row_groups_per_file : int, optional + Row groups per chunk file when ``use_row_groups`` is ``True`` (default 400). + path_dega_files : str or Path, optional + Landscape files directory used to load the gene mapping when + ``gene_str_to_int_mapping`` is not supplied. + gene_str_to_int_mapping : mapping, optional + Explicit gene name -> integer mapping. Overrides ``path_dega_files``. + level, drop_no_call, min_quality + See :func:`iter_zarr_transcript_batches`. + + Returns + ------- + tuple + ``(tile_bounds, tile_grid_info, chunk_info)``. ``chunk_info`` is ``None`` when + ``use_row_groups`` is ``False``. + """ + output_path = Path(path_output_dir) + output_path.mkdir(parents=True, exist_ok=True) + + transformation_matrix = np.loadtxt(str(path_transformation_matrix)) + sparse_matrix = csr_matrix(transformation_matrix) + gene_mapping = _resolve_gene_mapping(gene_str_to_int_mapping, path_dega_files) + + tmp_root = output_path / "_tmp_trx_zarr_build" + shards_dir = tmp_root / "shards" + spill_dir = tmp_root / "spill" + shards_dir.mkdir(parents=True, exist_ok=True) + spill_dir.mkdir(parents=True, exist_ok=True) + + try: + # Pass 1: transform each grid-position batch to a shard and track the extent. + max_x = 0.0 + max_y = 0.0 + shard_idx = 0 + for batch in iter_zarr_transcript_batches( + path_trx_zarr, + level=level, + drop_no_call=drop_no_call, + min_quality=min_quality, + ): + batch = _apply_gene_mapping(batch, gene_mapping) + transformed = _transform_batch(batch, sparse_matrix, image_scale) + + extent = transformed.select( + [ + pl.col("transformed_x").max().alias("mx"), + pl.col("transformed_y").max().alias("my"), + ] + ).row(0) + max_x = max(max_x, float(extent[0])) + max_y = max(max_y, float(extent[1])) + + transformed.write_parquet(shards_dir / f"shard_{shard_idx:06d}.parquet") + shard_idx += 1 + + x_min, y_min = 0.0, 0.0 + x_max, y_max = max_x, max_y + n_tiles_x = int(np.ceil((x_max - x_min) / tile_size)) if x_max > x_min else 0 + n_tiles_y = int(np.ceil((y_max - y_min) / tile_size)) if y_max > y_min else 0 + + tile_grid_info = { + "tile_size": tile_size, + "num_tiles_x": n_tiles_x, + "num_tiles_y": n_tiles_y, + "x_min": float(x_min), + "x_max": float(x_max), + "y_min": float(y_min), + "y_max": float(y_max), + } + tile_bounds = {"x_min": x_min, "x_max": x_max, "y_min": y_min, "y_max": y_max} + + chunk_info = None + if shard_idx > 0 and n_tiles_x > 0 and n_tiles_y > 0: + # Pass 2: spill each shard into per-tile parts. + for i, shard_path in enumerate(sorted(shards_dir.glob("shard_*.parquet"))): + _spill_one_transform_shard_to_tiles( + shard_path, i, x_min, y_min, n_tiles_x, n_tiles_y, tile_size, spill_dir + ) + shutil.rmtree(shards_dir, ignore_errors=True) + + if use_row_groups: + chunk_info = _write_tiles_as_row_groups_streaming( + str(output_path), tile_grid_info, spill_dir, max_row_groups_per_file + ) + else: + _write_traditional_transcript_tiles_from_spill( + spill_dir, n_tiles_x, n_tiles_y, str(output_path) + ) + + return tile_bounds, tile_grid_info, chunk_info + finally: + if tmp_root.exists(): + shutil.rmtree(tmp_root, ignore_errors=True) diff --git a/tests/integration/test_xenium_zarr_integration.py b/tests/integration/test_xenium_zarr_integration.py new file mode 100644 index 000000000..6a1f9b1ef --- /dev/null +++ b/tests/integration/test_xenium_zarr_integration.py @@ -0,0 +1,89 @@ +"""Integration test for tiling transcripts from a real Xenium ``transcripts.zarr.zip``. + +This test is skipped unless a real Xenium output bundle is available locally. Point the +``CELLDEGA_XENIUM_PANCREAS_DIR`` environment variable at an extracted Xenium ``outs`` +directory that contains ``transcripts.zarr.zip``, for example:: + + # Download + extract the 10x Xenium Human Pancreas FFPE dataset (~9 GB): + # https://www.10xgenomics.com/datasets/ffpe-human-pancreas-with-xenium-multimodal-cell-segmentation-1-standard + curl -O https://cf.10xgenomics.com/samples/xenium/2.0.0/Xenium_V1_human_Pancreas_FFPE/Xenium_V1_human_Pancreas_FFPE_outs.zip + unzip Xenium_V1_human_Pancreas_FFPE_outs.zip -d Xenium_V1_human_Pancreas_FFPE_outs + + export CELLDEGA_XENIUM_PANCREAS_DIR="$PWD/Xenium_V1_human_Pancreas_FFPE_outs" + pytest tests/integration/test_xenium_zarr_integration.py -q +""" + +import os +from pathlib import Path + +import numpy as np +import pytest + + +DATA_DIR = os.environ.get("CELLDEGA_XENIUM_PANCREAS_DIR") + + +def _zarr_bundle() -> Path | None: + if not DATA_DIR: + return None + candidate = Path(DATA_DIR) / "transcripts.zarr.zip" + return candidate if candidate.exists() else None + + +pytestmark = pytest.mark.skipif( + _zarr_bundle() is None, + reason=( + "Set CELLDEGA_XENIUM_PANCREAS_DIR to an extracted Xenium outs directory " + "containing transcripts.zarr.zip to run the integration test." + ), +) + + +def test_zarr_metadata_and_tiling(tmp_path): + from celldega.pre import trx_zarr + + bundle = _zarr_bundle() + + meta = trx_zarr.read_zarr_transcript_metadata(bundle) + assert meta["gene_names"], "expected a non-empty gene panel" + assert meta["number_rnas"] and meta["number_rnas"] > 0 + assert meta["position_keys"], "expected at least one level-0 grid position" + + # First batch should look like transcripts. + first_batch = next(trx_zarr.iter_zarr_transcript_batches(bundle)) + assert first_batch.columns == ["name", "x", "y"] + assert first_batch.height > 0 + + # Identity transform keeps coordinates in micron space for a self-contained test. + transform = tmp_path / "identity.csv" + np.savetxt(transform, np.eye(3)) + + gene_map = {name: idx for idx, name in enumerate(meta["gene_names"])} + out_dir = tmp_path / "transcripts" + + _tile_bounds, tile_grid_info, chunk_info = trx_zarr.make_trx_tiles_from_zarr( + bundle, + transform, + out_dir, + tile_size=200, + use_row_groups=True, + gene_str_to_int_mapping=gene_map, + ) + + assert tile_grid_info["num_tiles_x"] > 0 + assert tile_grid_info["num_tiles_y"] > 0 + assert chunk_info["total_row_groups"] == ( + tile_grid_info["num_tiles_x"] * tile_grid_info["num_tiles_y"] + ) + + # Total tiled transcripts should be positive and not exceed the raw count + # (no-call / low-quality transcripts are filtered out). + import pyarrow.parquet as pq + + total = 0 + for chunk_path in out_dir.glob("chunk_*.parquet"): + pf = pq.ParquetFile(chunk_path) + for rg in range(pf.metadata.num_row_groups): + total += pf.read_row_group(rg).num_rows + + assert 0 < total <= meta["number_rnas"] diff --git a/tests/unit/test_pre/test_trx_zarr.py b/tests/unit/test_pre/test_trx_zarr.py new file mode 100644 index 000000000..2aacd5f44 --- /dev/null +++ b/tests/unit/test_pre/test_trx_zarr.py @@ -0,0 +1,346 @@ +"""Tests for reading and tiling transcripts from the 10x ``transcripts.zarr.zip`` format. + +These cover: + +* metadata + batch reading of the ``/grids`` spatial pyramid, +* the no-call / quality filters, +* streaming tiling directly from the Zarr grid (traditional + row-group layouts), and +* parity between the Zarr tiling path and the flat-parquet ``make_trx_tiles`` path. +""" + +import importlib.util +import json +from pathlib import Path +import sys +import types + +import numpy as np +import pandas as pd +import pytest + + +try: + import polars as pl + import pyarrow.parquet as pq + import zarr +except (ImportError, ModuleNotFoundError) as e: # pragma: no cover - skip if deps missing + pytest.skip(f"Required libraries missing: {e}", allow_module_level=True) + + +# Dynamically load the pre submodules to avoid importing the heavy celldega.pre package. +ROOT_DIR = Path(__file__).resolve().parents[3] +PRE_ROOT = ROOT_DIR / "src" / "celldega" / "pre" +CELLPKG = types.ModuleType("celldega") +CELLPKG.__path__ = [str(ROOT_DIR / "src" / "celldega")] +sys.modules.setdefault("celldega", CELLPKG) +PREPKG = types.ModuleType("celldega.pre") +PREPKG.__path__ = [str(PRE_ROOT)] +sys.modules.setdefault("celldega.pre", PREPKG) + + +def _load(name): + spec = importlib.util.spec_from_file_location(f"celldega.pre.{name}", PRE_ROOT / f"{name}.py") + module = importlib.util.module_from_spec(spec) + module.__package__ = "celldega.pre" + sys.modules[f"celldega.pre.{name}"] = module + spec.loader.exec_module(module) + return module + + +boundary_tile = _load("boundary_tile") +trx_tile = _load("trx_tile") +trx_zarr = _load("trx_zarr") + +make_trx_tiles = trx_tile.make_trx_tiles + + +GENE_NAMES = ["G0", "G1", "G2"] +GENE_MAP = {name: idx for idx, name in enumerate(GENE_NAMES)} +NO_CALL = trx_zarr.NO_CALL_GENE_INDEX + + +def _write_zarr_transcripts(path, positions, *, gene_names=GENE_NAMES, as_zip=False): + """Write a synthetic ``transcripts.zarr`` mirroring the 10x /grids schema. + + ``positions`` maps a ``"col,row"`` key to a dict with ``x``, ``y``, ``gene`` (uint16 + gene indices) and optional ``quality`` lists. + """ + store = zarr.ZipStore(str(path), mode="w") if as_zip else zarr.DirectoryStore(str(path)) + root = zarr.group(store=store, overwrite=True) + total = sum(len(p["x"]) for p in positions.values()) + root.attrs["gene_names"] = list(gene_names) + root.attrs["number_rnas"] = total + root.attrs["spatial_units"] = "micron" + + grids = root.create_group("grids") + grids.attrs["number_levels"] = 1 + grids.attrs["grid_size"] = [100.0] + grids.attrs["grid_keys"] = [list(positions.keys())] + grids.attrs["grid_key_names"] = ["grid_x_loc", "grid_y_loc"] + + level0 = grids.create_group("0") + for key, data in positions.items(): + group = level0.create_group(key) + n = len(data["x"]) + location = np.column_stack( + [np.asarray(data["x"]), np.asarray(data["y"]), np.zeros(n)] + ).astype("float32") + group.create_dataset("location", data=location) + group.create_dataset("gene_identity", data=np.asarray(data["gene"], dtype="uint16")) + quality = data.get("quality", [40.0] * n) + group.create_dataset("quality_score", data=np.asarray(quality, dtype="float32")) + + if as_zip: + store.close() + return path + + +@pytest.fixture +def simple_zarr(tmp_path): + positions = { + "0,0": {"x": [50.0, 120.0], "y": [50.0, 120.0], "gene": [0, 1]}, + "1,1": {"x": [300.0, 320.0], "y": [60.0, 320.0], "gene": [2, 0]}, + } + return _write_zarr_transcripts(tmp_path / "transcripts.zarr", positions) + + +def test_is_zarr_transcript_path(): + assert trx_zarr.is_zarr_transcript_path("foo/transcripts.zarr.zip") + assert trx_zarr.is_zarr_transcript_path("foo/transcripts.zarr") + assert not trx_zarr.is_zarr_transcript_path("foo/transcripts.parquet") + + +def test_read_metadata(simple_zarr): + meta = trx_zarr.read_zarr_transcript_metadata(simple_zarr) + assert meta["gene_names"] == GENE_NAMES + assert meta["number_rnas"] == 4 + assert meta["number_levels"] == 1 + assert set(meta["position_keys"]) == {"0,0", "1,1"} + + +def test_iter_batches_yields_per_position(simple_zarr): + batches = list(trx_zarr.iter_zarr_transcript_batches(simple_zarr)) + assert len(batches) == 2 + for batch in batches: + assert batch.columns == ["name", "x", "y"] + assert batch.height == 2 + + combined = pl.concat(batches) + assert sorted(combined["name"].to_list()) == ["G0", "G0", "G1", "G2"] + + +def test_load_full_frame(simple_zarr): + frame = trx_zarr.load_zarr_transcripts(simple_zarr) + assert frame.height == 4 + assert set(frame["name"].to_list()) == {"G0", "G1", "G2"} + + +def test_drop_no_call(tmp_path): + positions = { + "0,0": {"x": [10.0, 20.0, 30.0], "y": [10.0, 20.0, 30.0], "gene": [0, NO_CALL, 2]}, + } + path = _write_zarr_transcripts(tmp_path / "transcripts.zarr", positions) + + kept = trx_zarr.load_zarr_transcripts(path, drop_no_call=True) + assert kept.height == 2 + assert sorted(kept["name"].to_list()) == ["G0", "G2"] + + all_rows = trx_zarr.load_zarr_transcripts(path, drop_no_call=False) + # The no-call index is out of range for gene_names, so it is filtered regardless. + assert all_rows.height == 2 + + +def test_min_quality_filter(tmp_path): + positions = { + "0,0": { + "x": [10.0, 20.0, 30.0], + "y": [10.0, 20.0, 30.0], + "gene": [0, 1, 2], + "quality": [10.0, 25.0, 40.0], + }, + } + path = _write_zarr_transcripts(tmp_path / "transcripts.zarr", positions) + + kept = trx_zarr.load_zarr_transcripts(path, min_quality=20.0) + assert sorted(kept["name"].to_list()) == ["G1", "G2"] + + +def test_zip_roundtrip(tmp_path): + positions = {"0,0": {"x": [1.0, 2.0], "y": [3.0, 4.0], "gene": [0, 2]}} + path = _write_zarr_transcripts(tmp_path / "transcripts.zarr.zip", positions, as_zip=True) + + meta = trx_zarr.read_zarr_transcript_metadata(path) + assert meta["number_rnas"] == 2 + frame = trx_zarr.load_zarr_transcripts(path) + assert sorted(frame["name"].to_list()) == ["G0", "G2"] + + +def _identity_transform(tmp_path): + path = tmp_path / "micron_to_image_transform.csv" + np.savetxt(path, np.eye(3)) + return path + + +def test_make_trx_tiles_from_zarr_traditional(tmp_path, simple_zarr): + transform = _identity_transform(tmp_path) + out_dir = tmp_path / "transcript_tiles" + tile_size = 250 + + tile_bounds, _tile_grid_info, chunk_info = trx_zarr.make_trx_tiles_from_zarr( + simple_zarr, + transform, + out_dir, + tile_size=tile_size, + use_row_groups=False, + gene_str_to_int_mapping=GENE_MAP, + ) + + assert chunk_info is None + assert tile_bounds["x_min"] == 0.0 + assert tile_bounds["x_max"] >= 320.0 + + tile_files = list(out_dir.glob("transcripts_tile_*.parquet")) + assert tile_files, "expected transcript tile files" + + total = 0 + for path in tile_files: + df = pd.read_parquet(path) + assert list(df.columns) == ["name", "geometry"] or set(df.columns) >= {"name", "geometry"} + # Gene names must be mapped to their integer codes. + assert set(df["name"].unique()) <= set(GENE_MAP.values()) + total += len(df) + assert total == 4 + + +def test_make_trx_tiles_from_zarr_row_groups(tmp_path, simple_zarr): + transform = _identity_transform(tmp_path) + out_dir = tmp_path / "transcripts" + tile_size = 250 + + _tile_bounds, tile_grid_info, chunk_info = trx_zarr.make_trx_tiles_from_zarr( + simple_zarr, + transform, + out_dir, + tile_size=tile_size, + use_row_groups=True, + max_row_groups_per_file=400, + gene_str_to_int_mapping=GENE_MAP, + ) + + assert chunk_info is not None + n_tiles = tile_grid_info["num_tiles_x"] * tile_grid_info["num_tiles_y"] + assert chunk_info["total_row_groups"] == n_tiles + + chunk_files = sorted(out_dir.glob("chunk_*.parquet")) + assert chunk_files + + # Metadata identifies the chunked row-group storage mode and grid dimensions. + pf = pq.ParquetFile(chunk_files[0]) + md = pf.schema_arrow.metadata + assert md[b"storage_mode"] == b"row_groups_chunked" + grid = json.loads(md[b"tile_grid_info"]) + assert grid["num_tiles_x"] == tile_grid_info["num_tiles_x"] + + # All transcripts are recoverable across the row groups. + total = 0 + for path in chunk_files: + file = pq.ParquetFile(path) + for rg in range(file.metadata.num_row_groups): + total += file.read_row_group(rg).num_rows + assert total == 4 + + +def test_empty_zarr_produces_no_tiles(tmp_path): + positions = {"0,0": {"x": [], "y": [], "gene": []}} + path = _write_zarr_transcripts(tmp_path / "transcripts.zarr", positions) + transform = _identity_transform(tmp_path) + out_dir = tmp_path / "transcript_tiles" + + _tile_bounds, _tile_grid_info, chunk_info = trx_zarr.make_trx_tiles_from_zarr( + path, transform, out_dir, tile_size=250, gene_str_to_int_mapping=GENE_MAP + ) + + assert chunk_info is None + assert list(out_dir.glob("transcripts_tile_*.parquet")) == [] + + +def test_zarr_matches_parquet_tiling(tmp_path): + """The Zarr grid tiler and the flat-parquet tiler must agree tile-for-tile.""" + # Interior coordinates (well away from 250 um tile boundaries) so both the + # floor-based Zarr assignment and the coarse/fine parquet filters agree exactly. + points = [ + (50.0, 50.0, "G0"), + (120.0, 120.0, "G1"), + (300.0, 60.0, "G2"), + (320.0, 320.0, "G0"), + (60.0, 300.0, "G1"), + (420.0, 420.0, "G2"), + ] + tile_size = 250 + + # meta_gene.parquet drives the gene->int mapping for the parquet path. + pd.DataFrame(index=GENE_NAMES).to_parquet(tmp_path / "meta_gene.parquet") + transform = _identity_transform(tmp_path) + + # Flat Xenium-style parquet source. + trx_parquet = tmp_path / "transcripts.parquet" + pl.DataFrame( + { + "feature_name": [g for _, _, g in points], + "x_location": [x for x, _, _ in points], + "y_location": [y for _, y, _ in points], + "cell_id": [f"c{i}" for i in range(len(points))], + "transcript_id": list(range(len(points))), + } + ).write_parquet(trx_parquet) + + # Equivalent Zarr source (two grid positions). + positions = { + "0,0": { + "x": [p[0] for p in points[:3]], + "y": [p[1] for p in points[:3]], + "gene": [GENE_MAP[p[2]] for p in points[:3]], + }, + "1,1": { + "x": [p[0] for p in points[3:]], + "y": [p[1] for p in points[3:]], + "gene": [GENE_MAP[p[2]] for p in points[3:]], + }, + } + zarr_path = _write_zarr_transcripts(tmp_path / "transcripts.zarr", positions) + + parquet_out = tmp_path / "transcript_tiles" + make_trx_tiles( + technology="Xenium", + path_trx=str(trx_parquet), + path_transformation_matrix=str(transform), + path_trx_tiles=str(parquet_out), + coarse_tile_factor=2, + tile_size=tile_size, + chunk_size=50, + image_scale=1, + max_workers=1, + ) + + zarr_out = tmp_path / "transcript_tiles_zarr" + trx_zarr.make_trx_tiles_from_zarr( + zarr_path, + transform, + zarr_out, + tile_size=tile_size, + use_row_groups=False, + path_dega_files=tmp_path, + ) + + def _tile_counts(directory): + counts = {} + for path in directory.glob("transcripts_tile_*.parquet"): + key = tuple(map(int, path.stem.split("_")[-2:])) + counts[key] = len(pd.read_parquet(path)) + return counts + + parquet_counts = _tile_counts(parquet_out) + zarr_counts = _tile_counts(zarr_out) + + assert parquet_counts == zarr_counts + assert sum(zarr_counts.values()) == len(points)