Skip to content

supprt zarr for xenium trx pre-processing #321

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
cornhundred wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

supprt zarr for xenium trx pre-processing #321

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 14 additions & 0 deletions src/celldega/pre/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down Expand Up @@ -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",
]
88 changes: 60 additions & 28 deletions src/celldega/pre/run_pre_processing.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down Expand Up @@ -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")
Expand All @@ -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")
Expand Down
160 changes: 90 additions & 70 deletions src/celldega/pre/trx_tile.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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"),
Expand Down Expand Up @@ -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)
Expand All @@ -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(
[
Expand Down Expand Up @@ -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():
Expand Down Expand Up @@ -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...")

Expand Down
Loading
Loading