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)