diff --git a/.gitignore b/.gitignore index 21e5524..4a404ab 100644 --- a/.gitignore +++ b/.gitignore @@ -24,3 +24,11 @@ xenium_example/ # docs /docs/generated/ /docs/_build/ + +# Local scratch work: new notebooks and scripts stay untracked by default +# (already-tracked files are unaffected; use `git add -f` to version a new one) +/docs/notebooks/*.ipynb +/scripts/*.py +/graphify-out/ +/slurm/ +/slurm-logs/ diff --git a/pyproject.toml b/pyproject.toml index d04033b..e8692d4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -149,7 +149,16 @@ lint.pydocstyle.convention = "numpy" files = [ "src", "tests" ] [[tool.mypy.overrides]] -module = [ "bioio.*", "geopandas.*", "shapely.*", "spatialdata.*", "spatialdata_io.*" ] +module = [ + "bioio.*", + "dask_image.*", + "geopandas.*", + "pyarrow.*", + "scipy.*", + "shapely.*", + "spatialdata.*", + "spatialdata_io.*", +] ignore_missing_imports = true [tool.pytest] diff --git a/src/spatialrefinery/core/downloader.py b/src/spatialrefinery/core/downloader.py index 9815bde..cfba899 100644 --- a/src/spatialrefinery/core/downloader.py +++ b/src/spatialrefinery/core/downloader.py @@ -3,7 +3,11 @@ `download_with_retries` is the reusable engine: it retries transient network failures and writes atomically (to `.part`, then `os.replace`s onto `dest`) so a killed download never leaves a truncated -file that a later run mistakes for complete. `BaseDownloader` wraps that +file that a later run mistakes for complete. A surviving `.part` is +resumed via an HTTP `Range` request rather than refetched from byte 0, +which matters for the multi-GB assets this package targets, and every +download whose length the server declares is size-checked before it is +promoted onto `dest`. `BaseDownloader` wraps that engine with a thread pool and turns per-asset results into a plan/run workflow that subclasses only need to feed with `RemoteAsset`s (see `spatialrefinery.io.xenium.XeniumDownloader`). @@ -21,6 +25,7 @@ from dataclasses import dataclass from pathlib import Path from typing import ClassVar, Literal +from urllib.error import HTTPError from urllib.request import Request, urlopen import spatialrefinery @@ -73,6 +78,40 @@ def ok(self) -> bool: return self.status in ("downloaded", "cached") +def _partial_size(part_path: Path) -> int: + """Bytes already fetched into `part_path`, or 0 if it is absent/unreadable.""" + try: + return part_path.stat().st_size + except OSError: + return 0 + + +def _declared_total(resp: object, headers: object) -> int | None: + """Total size of the *whole* resource, or None if the server didn't say. + + For a `206 Partial Content` reply the body is only the requested tail, + so `Content-Length` describes the tail, not the resource -- the total + lives in `Content-Range: bytes -/`. + """ + get = headers.get # type: ignore[attr-defined] + if getattr(resp, "status", None) == 206: + total = get("Content-Range", "").rpartition("/")[2].strip() + else: + total = (get("Content-Length") or "").strip() + return int(total) if total.isdigit() else None + + +def _content_length(url: str, headers: dict[str, str], timeout: int, ctx: ssl.SSLContext) -> int | None: + """Resource size via a HEAD request, or None if the server won't say.""" + try: + req = Request(url, headers=headers, method="HEAD") + with urlopen(req, timeout=timeout, context=ctx) as resp: + length = (resp.headers.get("Content-Length") or "").strip() + return int(length) if length.isdigit() else None + except OSError: + return None + + def download_with_retries( url: str, dest: str | Path, @@ -82,6 +121,7 @@ def download_with_retries( timeout: int = 180, chunk_size: int = DEFAULT_CHUNK_SIZE, overwrite: bool = False, + resume: bool = True, user_agent: str = DEFAULT_USER_AGENT, ) -> Path: """Stream `url` to `dest`, retrying transient failures, atomically. @@ -90,6 +130,14 @@ def download_with_retries( success, so an interrupted run never leaves a truncated file that a later run's `dest.exists()` check would mistake for complete. + With `resume=True` a `.part` left by an earlier attempt -- in this call + or in a previous process -- is continued with a `Range` request instead + of being refetched from byte 0. Servers that ignore `Range` answer + `200` with the whole body, which is detected and restarts the write, so + a resumed file is never a mix of two responses. Whenever the server + declares a length, the finished `.part` is size-checked before it is + promoted onto `dest`. + Parameters ---------- url : str @@ -105,11 +153,21 @@ def download_with_retries( overwrite : bool, optional If False (default) and `dest` already exists, return immediately without making a network request. + resume : bool, optional + Continue from a leftover `.part` using an HTTP `Range` + request, and keep that `.part` when an attempt fails so the next + attempt (or run) picks up where it stopped. Default True. Returns ------- Path `dest`, once fully and successfully written. + + Raises + ------ + RuntimeError + If every attempt fails, or the fetched byte count disagrees with + the size the server declared. """ dest = Path(dest) if dest.exists() and not overwrite: @@ -123,30 +181,71 @@ def download_with_retries( last_err: Exception | None = None for attempt in range(1, retries + 1): + offset = _partial_size(part_path) if resume else 0 try: - req = Request(url, headers=headers) - with urlopen(req, timeout=timeout, context=ctx) as resp, open(part_path, "wb") as out: - while chunk := resp.read(chunk_size): - out.write(chunk) + req_headers = dict(headers) + if offset: + req_headers["Range"] = f"bytes={offset}-" + req = Request(url, headers=req_headers) + with urlopen(req, timeout=timeout, context=ctx) as resp: + # A server free to ignore `Range` answers 200 with the entire + # body; appending that to what we already have would corrupt + # the file, so only 206 may append. + appending = bool(offset) and getattr(resp, "status", None) == 206 + if not appending: + offset = 0 + total = _declared_total(resp, resp.headers) + if appending: + logger.info("Resuming %s at %d bytes", dest.name, offset) + with open(part_path, "ab" if appending else "wb") as out: + while chunk := resp.read(chunk_size): + out.write(chunk) + + written = _partial_size(part_path) + if total is not None and written != total: + # Too long means the `.part` is not a prefix of the resource + # (a stale or corrupt leftover); resuming can never repair + # that, so drop it and let the next attempt start clean. + if written > total: + part_path.unlink(missing_ok=True) + raise OSError(f"Incomplete download: got {written} bytes, server declared {total}") + os.replace(part_path, dest) return dest + except HTTPError as e: + last_err = e + if e.code == 416 and offset: + # The range is past the end of the resource: either the file + # finished and only the rename was lost, or the `.part` is + # stale. Only the exact-size case is safe to promote. + total = _content_length(url, headers, timeout, ctx) + if total is not None and _partial_size(part_path) == total: + os.replace(part_path, dest) + return dest + part_path.unlink(missing_ok=True) + _backoff(attempt, retries, backoff, url, e) except OSError as e: # HTTPError, URLError, TimeoutError, and ssl.SSLError are all OSError # subclasses; catching OSError directly also cleans up `.part` for a # plain local write failure (e.g. disk full), which the previous - # narrower tuple missed. + # narrower tuple missed. With `resume` the bytes already on disk are + # a valid prefix, so they are kept for the next attempt. last_err = e - part_path.unlink(missing_ok=True) - if attempt < retries: - sleep_for = backoff ** (attempt - 1) - logger.warning( - "Retry %d/%d failed for %s: %s. Retrying in %.1fs...", attempt, retries, url, e, sleep_for - ) - time.sleep(sleep_for) + if not resume: + part_path.unlink(missing_ok=True) + _backoff(attempt, retries, backoff, url, e) raise RuntimeError(f"Failed to download after {retries} attempts: {url}\nLast error: {last_err}") +def _backoff(attempt: int, retries: int, backoff: float, url: str, err: Exception) -> None: + """Sleep between attempts, unless `attempt` was the last one.""" + if attempt < retries: + sleep_for = backoff ** (attempt - 1) + logger.warning("Retry %d/%d failed for %s: %s. Retrying in %.1fs...", attempt, retries, url, err, sleep_for) + time.sleep(sleep_for) + + class BaseDownloader(ABC): """Fetch the asset bundle for one spatial-omics technology. @@ -168,6 +267,7 @@ def __init__( timeout: int = 180, extract: bool = True, overwrite: bool = False, + resume: bool = True, dry_run: bool = False, ) -> None: self.outdir = Path(outdir) @@ -177,6 +277,7 @@ def __init__( self.timeout = timeout self.extract = extract self.overwrite = overwrite + self.resume = resume self.dry_run = dry_run # ---- must implement --------------------------------------------- # @@ -216,6 +317,7 @@ def fetch(self, asset: RemoteAsset) -> DownloadResult: backoff=self.backoff, timeout=self.timeout, overwrite=self.overwrite, + resume=self.resume, ) except Exception as e: # noqa: BLE001 - one asset's failure must not abort the whole batch return DownloadResult(asset=asset, path=None, status="failed", error=str(e)) diff --git a/src/spatialrefinery/core/utils.py b/src/spatialrefinery/core/utils.py index 87b9264..9ab0655 100644 --- a/src/spatialrefinery/core/utils.py +++ b/src/spatialrefinery/core/utils.py @@ -167,12 +167,20 @@ def bin_centroids(x: np.ndarray, y: np.ndarray, spot_size_um: float) -> tuple[np return centroid_x, centroid_y -def hex_grid_centroids( +def hex_lattice_params( bounds: tuple[float, float, float, float], spot_size_um: float, overlap: float | None = 0.0, -) -> np.ndarray: - """Generate a pointy-top hexagonal grid of centroids covering `bounds`. +) -> dict: + """Describe the pointy-top hexagonal lattice that covers `bounds`. + + This is the single definition of the lattice: :func:`hex_grid_centroids` + materialises it as centroids, and :func:`assign_points_to_hexes` uses it to + decide membership analytically. Keeping both on top of this function is what + guarantees the two can never drift apart. + + Centroid `k` of the grid sits at row `i = k // nx`, column `j = k % nx`, with + odd rows shifted right by `dx / 2`. Parameters ---------- @@ -187,16 +195,16 @@ def hex_grid_centroids( Returns ------- - np.ndarray - Array of shape `(N, 2)` with one `(x, y)` centroid per row. Empty - (`shape (0, 2)`) if the grid has no cells. + dict + `s` (circumradius), `dx`/`dy` (centroid spacing), `x_range`/`y_range` + (row/column origins) and `nx`/`ny` (grid shape). """ overlap = overlap or 0.0 if not (0.0 <= overlap < 1.0): raise ValueError(f"overlap must be in [0.0, 1.0), found {overlap}") x_min, y_min, x_max, y_max = bounds - s = spot_size_um / 2.0 # side length + s = spot_size_um / 2.0 # side length == circumradius hex_width = np.sqrt(3) * s hex_height = 2 * s @@ -206,15 +214,52 @@ def hex_grid_centroids( x_range = np.arange(x_min - dx, x_max + dx, dx) y_range = np.arange(y_min - dy, y_max + dy, dy) - centroids = [] - for i, y_val in enumerate(y_range): - x_offset = dx / 2 if i % 2 else 0 - for x_val in x_range: - centroids.append((x_val + x_offset, y_val)) + return { + "s": s, + "dx": dx, + "dy": dy, + "x_range": x_range, + "y_range": y_range, + "nx": len(x_range), + "ny": len(y_range), + } + + +def hex_grid_centroids( + bounds: tuple[float, float, float, float], + spot_size_um: float, + overlap: float | None = 0.0, +) -> np.ndarray: + """Generate a pointy-top hexagonal grid of centroids covering `bounds`. + + Parameters + ---------- + bounds : tuple[float, float, float, float] + `(x_min, y_min, x_max, y_max)`, as returned by :func:`xy_bounds`. + spot_size_um : float + Hexagon diameter (distance between opposite vertices), in the same + units as `bounds`. + overlap : float, optional + Fractional overlap between adjacent hexagons, applied by reducing + centroid spacing while keeping hexagon size constant. Default 0.0. + + Returns + ------- + np.ndarray + Array of shape `(N, 2)` with one `(x, y)` centroid per row, ordered + row-major (all of row 0 left to right, then row 1, ...). Empty + (`shape (0, 2)`) if the grid has no cells. + """ + lattice = hex_lattice_params(bounds, spot_size_um, overlap) + x_range, y_range = lattice["x_range"], lattice["y_range"] + nx, ny = lattice["nx"], lattice["ny"] - if not centroids: + if nx == 0 or ny == 0: return np.empty((0, 2)) - return np.asarray(centroids) + + rows = np.repeat(np.arange(ny), nx) + x_offset = np.where(rows % 2 == 1, lattice["dx"] / 2, 0.0) + return np.column_stack([np.tile(x_range, ny) + x_offset, np.repeat(y_range, nx)]) def hexagon_vertices(center_x: float, center_y: float, side_length: float) -> np.ndarray: @@ -225,6 +270,203 @@ def hexagon_vertices(center_x: float, center_y: float, side_length: float) -> np return np.column_stack([x, y]) +def hex_candidate_offsets(overlap: float | None = 0.0) -> tuple[range, range]: + """Return the `(row, column)` lattice offsets a point could fall into. + + A point at lattice cell `(i0, j0)` can only lie inside hexagons whose + centroid is within one circumradius `s` of it. With centroid spacing + `dy = 1.5 * s * (1 - overlap)` and `dx = sqrt(3) * s * (1 - overlap)`, that + is always rows `{i0, i0 + 1}` and columns `{j0, j0 + 1}`; wider overlaps + bring further neighbours into reach, so the span is derived from `overlap` + rather than hard-coded (the `0.06` used in practice needs only the base + 2x2 block). + + Returns + ------- + tuple[range, range] + `(row_offsets, column_offsets)` to try around the containing cell. + """ + overlap = overlap or 0.0 + if not (0.0 <= overlap < 1.0): + raise ValueError(f"overlap must be in [0.0, 1.0), found {overlap}") + + # s / dy and (sqrt(3)/2 * s) / dx respectively, i.e. how many extra whole + # centroid spacings fit inside the hexagon's reach. + extra_rows = int(np.floor(1.0 / (1.5 * (1 - overlap)))) + extra_cols = int(np.floor(1.0 / (2.0 * (1 - overlap)))) + return range(-extra_rows, 2 + extra_rows), range(-extra_cols, 2 + extra_cols) + + +def assign_points_to_hexes( + x: np.ndarray, + y: np.ndarray, + lattice: dict, + overlap: float | None = 0.0, +) -> tuple[np.ndarray, np.ndarray]: + """Map points to every lattice hexagon that contains them. + + The lattice is regular, so membership is a closed-form test rather than a + spatial join: no geometry objects are built and nothing is indexed. A point + on a shared edge is reported for both hexagons, matching the `intersects` + predicate `GeoDataFrame.sjoin` uses. When `overlap > 0` the hexagons + genuinely overlap, so a point can be returned for more than one spot -- the + same duplication the sjoin produces. + + Parameters + ---------- + x, y : np.ndarray + Point coordinates, in the units `lattice` was built in. + lattice : dict + As returned by :func:`hex_lattice_params`. + overlap : float, optional + The overlap `lattice` was built with; sets how far to search. + + Returns + ------- + tuple[np.ndarray, np.ndarray] + `(point_positions, spot_ids)`, where `spot_ids` index + :func:`hex_grid_centroids` output for the same lattice. + """ + x = np.asarray(x, dtype=np.float64) + y = np.asarray(y, dtype=np.float64) + + s, dx, dy = lattice["s"], lattice["dx"], lattice["dy"] + nx, ny = lattice["nx"], lattice["ny"] + if nx == 0 or ny == 0 or x.size == 0: + return np.empty(0, dtype=np.int64), np.empty(0, dtype=np.int64) + + x0, y0 = lattice["x_range"][0], lattice["y_range"][0] + half_width = np.sqrt(3) / 2 * s + inv_sqrt3 = 1.0 / np.sqrt(3) + + row_offsets, col_offsets = hex_candidate_offsets(overlap) + base_row = np.floor((y - y0) / dy).astype(np.int64) + + positions: list[np.ndarray] = [] + spot_ids: list[np.ndarray] = [] + for row_offset in row_offsets: + row = base_row + row_offset + row_ok = (row >= 0) & (row < ny) + if not row_ok.any(): + continue + centre_y = y0 + row * dy + abs_dy = np.abs(y - centre_y) + # Odd rows are shifted right by half a spacing, so the column origin + # depends on the row -- recompute it rather than reusing base_row's. + x_offset = np.where(row % 2 == 1, dx / 2, 0.0) + base_col = np.floor((x - x0 - x_offset) / dx).astype(np.int64) + for col_offset in col_offsets: + col = base_col + col_offset + centre_x = x0 + x_offset + col * dx + abs_dx = np.abs(x - centre_x) + inside = row_ok & (col >= 0) & (col < nx) & (abs_dx <= half_width) & (abs_dy <= s - abs_dx * inv_sqrt3) + if inside.any(): + hit = np.flatnonzero(inside) + positions.append(hit) + spot_ids.append(row[hit] * nx + col[hit]) + + if not positions: + return np.empty(0, dtype=np.int64), np.empty(0, dtype=np.int64) + return np.concatenate(positions), np.concatenate(spot_ids) + + +def _merge_pair_counts(pending, keys, counts): + """Fold a list of raw pair keys into running `(unique keys, counts)` totals.""" + stacked = [*pending] + weights = [np.ones(part.size, dtype=np.int64) for part in pending] + if keys.size: + stacked.append(keys) + weights.append(counts) + + all_keys = np.concatenate(stacked) + if all_keys.size == 0: + return keys, counts + all_weights = np.concatenate(weights) + + order = np.argsort(all_keys, kind="stable") + all_keys = all_keys[order] + all_weights = all_weights[order] + unique_keys, starts = np.unique(all_keys, return_index=True) + return unique_keys, np.add.reduceat(all_weights, starts) + + +def bin_points_to_hex_counts( + batches, + lattice: dict, + n_spots: int, + n_genes: int, + overlap: float | None = 0.0, + flush_pairs: int = 256_000_000, +): + """Accumulate per-(spot, gene) counts by streaming batches of points. + + Peak memory is one batch plus the running set of non-zero cells, so this + scales to billion-transcript sections that cannot be held in memory at + once. `spatialdata.aggregate` cannot: it builds one shapely `Point` per + transcript and then groups with `observed=False`, which materialises the + full `n_spots x n_genes` product regardless of how sparse the data is. + + Parameters + ---------- + batches : Iterable[tuple[np.ndarray, np.ndarray, np.ndarray]] + `(x, y, gene_codes)` triples. `gene_codes` index the gene axis; + negative codes are dropped, which is how unmapped features are skipped. + lattice : dict + As returned by :func:`hex_lattice_params`. + n_spots, n_genes : int + Shape of the matrix to build. + overlap : float, optional + The overlap `lattice` was built with. + flush_pairs : int, optional + Merge pending `(spot, gene)` pairs into the running totals once this + many have accumulated. Bounds the merge buffer; does not change output. + Each merge re-sorts the running totals, so this trades peak memory + against the number of merges: a 2.7e9-transcript section flushes ~12 + times at the default and ~48 times at a quarter of it, with the running + array growing to hundreds of millions of cells either way. + + Returns + ------- + scipy.sparse.csr_matrix + Counts of shape `(n_spots, n_genes)`. + """ + from scipy import sparse + + # One integer key per (spot, gene) cell keeps each merge a single sort + # rather than a lexsort over two columns. + running_keys = np.empty(0, dtype=np.int64) + running_counts = np.empty(0, dtype=np.int64) + pending: list[np.ndarray] = [] + pending_size = 0 + + for batch_x, batch_y, gene_codes in batches: + positions, spot_ids = assign_points_to_hexes(batch_x, batch_y, lattice, overlap) + if positions.size == 0: + continue + codes = np.asarray(gene_codes)[positions].astype(np.int64) + keep = codes >= 0 + if not keep.all(): + spot_ids = spot_ids[keep] + codes = codes[keep] + if spot_ids.size == 0: + continue + + pending.append(spot_ids * n_genes + codes) + pending_size += spot_ids.size + if pending_size >= flush_pairs: + running_keys, running_counts = _merge_pair_counts(pending, running_keys, running_counts) + pending, pending_size = [], 0 + + if pending: + running_keys, running_counts = _merge_pair_counts(pending, running_keys, running_counts) + + if running_keys.size == 0: + return sparse.csr_matrix((n_spots, n_genes), dtype=np.int64) + + rows, cols = np.divmod(running_keys, n_genes) + return sparse.coo_matrix((running_counts, (rows, cols)), shape=(n_spots, n_genes), dtype=np.int64).tocsr() + + def affine_from_point_pairs(src_pts: np.ndarray, dst_pts: np.ndarray) -> np.ndarray: """Compute a 3x3 affine matrix mapping `src_pts` to `dst_pts` (3 point pairs).""" import cv2 @@ -265,6 +507,62 @@ def fix_table_validation_errors(adata): return adata +def decode_bytes_columns(points): + """Decode bytes-valued ``object`` columns of a points frame to pandas strings. + + Older Xenium outputs (the "Preview" / "With_Addon" bundles) store + `cell_id` and `fov_name` in `transcripts.parquet` as parquet *binary* + rather than *string*, so `spatialdata_io.xenium` hands back a dask frame + whose columns carry `bytes` under a bare `object` dtype. + + That combination cannot be written. `SpatialData.write` sends points + through `dask.dataframe.to_parquet`, which infers the Arrow schema from + `meta_nonempty(df._meta)`; for a bare `object` dtype dask fills the dummy + frame with a literal `object()` sentinel, and pyarrow raises + `ArrowInvalid: ... did not recognize Python value type` before any real + data is touched. Re-typing the columns as `string` makes both the dummy + frame and the real values inferable. + + Parameters + ---------- + points : dask.dataframe.DataFrame + A points element, typically `sdata["transcripts"]`. + + Returns + ------- + dask.dataframe.DataFrame + The same element with every `object` column re-typed to `string`, and + its coordinate transformations preserved. Returned unchanged when + there are no `object` columns. + """ + from spatialdata.models import PointsModel + from spatialdata.transformations import get_transformation, set_transformation + + object_columns = [name for name, dtype in points.dtypes.items() if pd.api.types.is_object_dtype(dtype)] + if not object_columns: + return points + + logger.info("Re-typing bytes/object columns to string: %s", object_columns) + + # get_all=True so a frame registered in several coordinate systems keeps + # every one of them, not just "global". + transformations = get_transformation(points, get_all=True) + + decoded = points + for column in object_columns: + decoded[column] = decoded[column].map( + lambda value: value.decode("utf-8") if isinstance(value, bytes) else value, + meta=(column, "string"), + ) + + # map() drops the PointsModel attrs, so re-parse and re-attach the + # transformations captured above. + decoded = PointsModel.parse(decoded) + set_transformation(decoded, transformations, set_all=True) + + return decoded + + def slide_to_numpy(slide, level: int = 0) -> np.ndarray: """Convert an `openslide.OpenSlide` object to a NumPy array. @@ -446,6 +744,7 @@ def create_hexagonal_spots( key_x: str = "x", key_y: str = "y", overlap: float | None = 0.0, + bounds: tuple[float, float, float, float] | None = None, ): """Cover the extent of `df` with hexagonal pseudo-spots. @@ -465,6 +764,11 @@ def create_hexagonal_spots( Fractional overlap between adjacent hexagons (e.g. 0.06 for 6% overlap), applied by reducing centroid spacing while keeping hexagon size constant. Default 0.0. + bounds : tuple[float, float, float, float], optional + Precomputed `(x_min, y_min, x_max, y_max)` for `df`. Pass this when the + caller already has the extent: deriving it from a lazy dask frame costs + a full scan, and callers that also need the lattice want both to come + from the same numbers. Returns ------- @@ -476,7 +780,8 @@ def create_hexagonal_spots( from shapely.geometry import Polygon overlap = overlap or 0.0 - bounds = xy_bounds(df, key_x, key_y) + if bounds is None: + bounds = xy_bounds(df, key_x, key_y) s = spot_size_um / 2.0 centroids = hex_grid_centroids(bounds, spot_size_um, overlap) diff --git a/src/spatialrefinery/io/xenium.py b/src/spatialrefinery/io/xenium.py index 2e020b6..374c1c0 100644 --- a/src/spatialrefinery/io/xenium.py +++ b/src/spatialrefinery/io/xenium.py @@ -17,10 +17,11 @@ from pathlib import Path from typing import Any +import numpy as np import pandas as pd import spatialdata as sd from PIL import Image -from spatialdata.models import ShapesModel +from spatialdata.models import ShapesModel, TableModel from spatialdata.transformations import Identity, Scale, get_transformation, set_transformation from spatialdata.transformations.transformations import Affine from spatialdata_io import xenium, xenium_aligned_image @@ -29,13 +30,17 @@ from spatialrefinery.core.downloader import BaseDownloader, DownloadResult, RemoteAsset from spatialrefinery.core.registry import TechnologySpec, register_technology from spatialrefinery.core.utils import ( + bin_points_to_hex_counts, create_circular_spots, create_hexagonal_spots, + decode_bytes_columns, fix_table_validation_errors, + hex_lattice_params, parse_curl_manifest, segment_tissue, slide_to_numpy, transform_name, + xy_bounds, ) logger = logging.getLogger(__name__) @@ -81,7 +86,192 @@ def read_xenium_alignment(alignment_file_path: str): return np.asarray(alignment_matrix) -def create_pseudo_spots(sdata, spot_size_um: int = 55, overlap: float | None = 0.06, values: str = "transcripts"): +def _aligned_image_dims(image_path: str) -> tuple[str, ...] | None: + """Return an explicit `dims` for `xenium_aligned_image`, or None to let it infer. + + `xenium_aligned_image` guesses the axis order from the array shape alone and + only handles 10x's interleaved layout, asserting `image.shape[-1] == 3` for + 4-D input. Scanners that write the OME-TIFF with `planarconfig=SEPARATE` + produce `(1, c, y, x)` instead, which trips that assert. Probing the lazy + array is cheap -- no pixels are read -- and the reader squeezes any axis not + named "c"/"x"/"y", so a dummy leading axis is the documented way through. + + Returns + ------- + tuple[str, ...] | None + `dims` to pass through, or `None` when the reader's own inference is + already correct. + """ + from dask_image.imread import imread + + shape = imread(image_path).shape + if len(shape) == 4: + if shape[-1] in (3, 4): + return None # (1, y, x, c) -- what the reader already expects + if shape[1] in (3, 4): + logger.info("H&E image is planar %s; parsing as (dummy, c, y, x)", shape) + return ("dummy", "c", "y", "x") + elif len(shape) == 3 and shape[0] in (3, 4): + return None # (c, y, x) -- also handled by the reader + + raise ValueError( + f"Cannot infer the channel axis of {image_path} from shape {shape}; " + "pass dims= explicitly to xenium_aligned_image." + ) + + +def _transcript_gene_index(points) -> pd.Index: + """Return the gene axis to aggregate transcripts onto. + + Uses the categorical's own categories, which is what `SpatialData.aggregate` + used for its columns, so the resulting table keeps the same gene order the + previous implementation produced. + + The subtlety is that dask marks a categorical it read lazily from parquet as + *unknown*: `dtype.categories` is then a single `'__UNKNOWN_CATEGORIES__'` + placeholder rather than the real gene list. Taking it at face value maps + every transcript to a code of -1 and silently yields an all-zero table, so + resolve it first. `sdata` comes back off disk that way in the conversion + pipeline, while a freshly-read `xenium()` object has known categories -- + hence both paths. + """ + feature = points["feature_name"] + dtype = getattr(feature, "dtype", None) + + if not isinstance(dtype, pd.CategoricalDtype): + values = feature.compute() if hasattr(feature, "compute") else feature + return pd.Index(sorted(pd.unique(values.dropna()))) + + try: + from dask.dataframe.utils import has_known_categories + + known = has_known_categories(feature) + except (ImportError, TypeError, AttributeError): + known = True + + if not known: + logger.info("Resolving transcripts' feature_name categories (one pass over that column)") + feature = feature.cat.as_known() + + genes = pd.Index(feature.dtype.categories) + if len(genes) == 0: + raise ValueError("transcripts' feature_name has no categories; cannot build a gene axis") + return genes + + +def _gene_codes(column, genes: pd.Index) -> np.ndarray: + """Map one batch of `feature_name` values onto positions in `genes` (-1 = drop).""" + if isinstance(column.dtype, pd.CategoricalDtype): + # Map the (small) category list once, then take through the codes -- + # far cheaper than hashing every one of a batch's values. + category_positions = genes.get_indexer(column.cat.categories) + raw = column.cat.codes.to_numpy() + codes = np.full(raw.shape, -1, dtype=np.int64) + present = raw >= 0 + codes[present] = category_positions[raw[present]] + return codes + return genes.get_indexer(column.to_numpy()).astype(np.int64) + + +def _aggregate_transcripts_hex(sdata, spots_name: str, lattice: dict, overlap: float | None): + """Count transcripts per pseudo-spot by streaming them over the hex lattice. + + Replaces `SpatialData.aggregate` for the transcripts-to-spots case. That + path converts every transcript into a shapely `Point` before sjoining, then + groups with `observed=False`, which materialises the whole + `n_spots x n_genes` product; on a whole-transcriptome section (18k targets, + ~1e9 transcripts) that needs far more than the 500 GB a node has. The + lattice is regular, so membership is closed-form: this streams the points + once and keeps only non-zero cells. + + Overlapping spots still double-count a transcript that falls in more than + one hexagon, exactly as the sjoin did. + + Parameters + ---------- + sdata : spatialdata.SpatialData + Must contain `"transcripts"` and the `spots_name` shapes element. + spots_name : str + Name of the pseudo-spot shapes element, e.g. `"spots_55um"`. + lattice : dict + The lattice `spots_name` was built from, per `hex_lattice_params`. + overlap : float, optional + The overlap that lattice was built with. + + Returns + ------- + anndata.AnnData + Counts for every spot in `sdata[spots_name]`, in that element's order. + """ + import anndata as ad + + points = sdata["transcripts"] + spots = sdata[spots_name] + genes = _transcript_gene_index(points) + + def batches(): + columns = ["x", "y", "feature_name"] + for partition in points.partitions: + frame = partition[columns].compute() + if frame.empty: + continue + yield ( + frame["x"].to_numpy(dtype=np.float64), + frame["y"].to_numpy(dtype=np.float64), + _gene_codes(frame["feature_name"], genes), + ) + + logger.info( + "Binning transcripts onto %d spots x %d genes (streaming %d partitions)", + len(spots), + len(genes), + points.npartitions, + ) + counts = bin_points_to_hex_counts(batches(), lattice, n_spots=len(spots), n_genes=len(genes), overlap=overlap) + if counts.nnz == 0: + # Every transcript missing every spot means the lattice and the points + # disagree, or the gene axis never matched -- not a real empty section. + # Writing the all-zero table that results is worse than stopping. + raise ValueError( + f"Aggregating transcripts onto {spots_name!r} produced no counts at all " + f"({len(spots)} spots, {len(genes)} genes). The spot lattice and the " + "transcript coordinates are inconsistent, or feature_name did not map " + "onto the gene axis." + ) + + logger.info( + "Binned transcripts into %d non-zero (spot, gene) cells (%.2f%% dense)", + counts.nnz, + 100.0 * counts.nnz / max(counts.shape[0] * counts.shape[1], 1), + ) + + # Parse as a TableModel, not a bare AnnData: `SpatialData.aggregate` returned + # a parsed table, and it is `uns["spatialdata_attrs"]` plus the region / + # instance_id columns it writes that register this table as *annotating* + # `spots_name`. Without them the counts are still correct but nothing links + # them to the spots, and consumers (spatialdata-plot among them) report the + # element as having no annotating tables. + table = ad.AnnData( + X=counts, + obs=pd.DataFrame( + { + "instance_id": spots.index.to_numpy(), + "region": pd.Categorical([spots_name] * len(spots), categories=[spots_name]), + }, + index=spots.index.astype(str), + ), + var=pd.DataFrame(index=genes.astype(str)), + ) + return TableModel.parse(table, region=spots_name, region_key="region", instance_key="instance_id") + + +def create_pseudo_spots( + sdata, + spot_size_um: int = 55, + overlap: float | None = 0.06, + values: str = "transcripts", + max_spots_per_chunk: int | None = 50_000, +): """ Create pseudo-spots from transcripts and add them to the SpatialData object. @@ -100,6 +290,12 @@ def create_pseudo_spots(sdata, spot_size_um: int = 55, overlap: float | None = 0 values : str, optional What to aggregate into each pseudo-spot: "transcripts" (default) or "cell_boundaries". + max_spots_per_chunk : int, optional + Retained for backwards compatibility and ignored. Transcript + aggregation now streams over the hex lattice (see + :func:`_aggregate_transcripts_hex`), so peak memory is bounded by one + dask partition plus the non-zero output cells and no longer depends on + how many spots there are. Returns ------- @@ -111,9 +307,15 @@ def create_pseudo_spots(sdata, spot_size_um: int = 55, overlap: float | None = 0 if values not in ("transcripts", "cell_boundaries"): raise ValueError(f"values can only be one of ['transcripts', 'cell_boundaries'], found {values!r}") + # Derive the extent once and build both the polygons and the lattice from + # it: the aggregation below bins against the same lattice the polygons come + # from, so they must not be computed from two separate dask scans. + bounds = xy_bounds(sdata["transcripts"], "x", "y") + lattice = hex_lattice_params(bounds, spot_size_um, overlap) + # Create the pseudo-spots hexagonal_spots_gdf = create_hexagonal_spots( - sdata["transcripts"], key_x="x", key_y="y", spot_size_um=spot_size_um, overlap=overlap + sdata["transcripts"], key_x="x", key_y="y", spot_size_um=spot_size_um, overlap=overlap, bounds=bounds ) # Create transformation for the spot polygons @@ -139,15 +341,11 @@ def create_pseudo_spots(sdata, spot_size_um: int = 55, overlap: float | None = 0 set_transformation(sdata["transcripts"], transcripts_2d_transform, to_coordinate_system="global") # Pool transcripts (or cell_boundaries) to pseudo-spots using the hexagonal polygons + spots_name = f"spots_{spot_size_um!s}um" if values == "transcripts": - aggr_sdata = sdata.aggregate( - values="transcripts", - by=f"spots_{spot_size_um!s}um", - value_key="feature_name", - agg_func="count", - table_name="table", - deepcopy=False, - ) + # Deliberately not SpatialData.aggregate: see _aggregate_transcripts_hex. + aggregated_table = _aggregate_transcripts_hex(sdata, spots_name, lattice, overlap) + aggr_sdata = None else: # values == "cell_boundaries" aggr_sdata = sdata.aggregate( values="cell_boundaries", @@ -158,30 +356,38 @@ def create_pseudo_spots(sdata, spot_size_um: int = 55, overlap: float | None = 0 table_name="table", deepcopy=False, ) + aggregated_table = aggr_sdata["table"] # Ensure only common genes are included - common_genes = aggr_sdata["table"].var.index.intersection(sdata["table"].var.index) + common_genes = aggregated_table.var.index.intersection(sdata["table"].var.index) + if len(common_genes) == 0: + raise ValueError( + f"No gene names shared between the aggregated spots ({len(aggregated_table.var)} features) " + f"and the cell table ({len(sdata['table'].var)} features); the spot table would be empty." + ) # Filter to common genes in-place (creates a view, not a full copy) - aggr_sdata["table"] = aggr_sdata["table"][:, common_genes] + aggregated_table = aggregated_table[:, common_genes] - # Add spatial coordinates directly to the table - aggr_sdata["table"].obsm["spatial"] = aggr_sdata[f"spots_{spot_size_um!s}um"][["x_um", "y_um"]].values + # Add spatial coordinates directly to the table. Read them off the spots + # element rather than the aggregation result: the tiled path has no + # aggr_sdata, and the table is ordered to match this element either way. + aggregated_table.obsm["spatial"] = sdata[spots_name][["x_um", "y_um"]].values # Determine which spots are "in tissue" by checking intersection with cell boundaries. # `tissue_contours` is only written when an aligned H&E image was found and tissue # segmentation succeeded (see xenium_to_spatialdata); without it, every spot's # in-tissue status is unknown, so default to 1 rather than raising. if "tissue_contours" in sdata: - spot_gdf = sd.transform(sdata[f"spots_{spot_size_um!s}um"], to_coordinate_system="global").copy() + spot_gdf = sd.transform(sdata[spots_name], to_coordinate_system="global").copy() contour_gdf = sd.transform(sdata["tissue_contours"], to_coordinate_system="global").copy() # Perform spatial join to find spots that intersect with tissue contours intersecting_polygons = gpd.sjoin(spot_gdf, contour_gdf, how="inner", predicate="intersects") # Intersect with obs.index first: `.loc` setitem with unknown labels would # silently *enlarge* obs with spurious rows instead of raising. - in_tissue_indices = intersecting_polygons.index.unique().astype(str).intersection(aggr_sdata["table"].obs.index) - aggr_sdata["table"].obs["in_tissue"] = aggr_sdata["table"].obs.index.isin(in_tissue_indices).astype(int) + in_tissue_indices = intersecting_polygons.index.unique().astype(str).intersection(aggregated_table.obs.index) + aggregated_table.obs["in_tissue"] = aggregated_table.obs.index.isin(in_tissue_indices).astype(int) del spot_gdf, contour_gdf, intersecting_polygons else: @@ -189,16 +395,16 @@ def create_pseudo_spots(sdata, spot_size_um: int = 55, overlap: float | None = 0 "No 'tissue_contours' found in sdata; marking all spots as in_tissue=1 " "(no H&E image was aligned, or tissue segmentation failed)." ) - aggr_sdata["table"].obs["in_tissue"] = 1 + aggregated_table.obs["in_tissue"] = 1 # Add to SpatialData object - sdata[f"spots_{spot_size_um!s}um_table"] = aggr_sdata["table"] + sdata[f"spots_{spot_size_um!s}um_table"] = aggregated_table logger.info("Updating sdata with table element for spots of size %sµm...", spot_size_um) sdata.write_element(f"spots_{spot_size_um!s}um_table") # Clean up aggregated sdata - del aggr_sdata + del aggr_sdata, aggregated_table return sdata @@ -217,9 +423,21 @@ def find_xenium_files(dataset_path: Path) -> dict: dict Dictionary containing paths to found files (None if not found). """ + # `*he_imagealignment.csv` is 10x's own name, but bundles from other + # pipelines (e.g. the Atera "WTA Preview" sets) ship + # `_he_alignment.csv` instead. Missing it is not harmless: the H&E + # then lands with an Identity transform, silently unaligned, and + # `tissue_contours` with it. Keypoint files are deliberately not matched + # here -- they are a different format, handled by `read_xenium_alignment` + # only when they are the alignment file itself. file_patterns = { "img_path": ["*_he_*.tiff", "*_he_*.tif", "*_he_*.svs", "*_he_*.ndpi"], - "alignment_file_path": ["*he_imagealignment.csv"], + "alignment_file_path": [ + "*he_imagealignment.csv", + "*_he_alignment.csv", + "*_imagealignment.csv", + "*_alignment.csv", + ], "experiment_path": ["experiment.xenium"], } @@ -227,9 +445,11 @@ def find_xenium_files(dataset_path: Path) -> dict: for file_key, patterns in file_patterns.items(): found_file = None for pattern in patterns: - matches = list(dataset_path.glob(pattern)) + # sorted() so a directory holding several matches resolves the same + # way on every filesystem, rather than following readdir order. + matches = sorted(dataset_path.glob(pattern)) if matches: - found_file = matches[0] # Take the first match + found_file = matches[0] break paths[file_key] = found_file @@ -247,6 +467,7 @@ def xenium_to_spatialdata( values: str = "transcripts", n_jobs: int = 1, overwrite: bool = False, + max_spots_per_chunk: int | None = 50_000, ) -> Path: """ Convert 10x Xenium raw data to SpatialData zarr format. @@ -283,6 +504,9 @@ def xenium_to_spatialdata( Number of workers for parallel processing. Default is 1. overwrite : bool, optional Whether to overwrite existing zarr file. Default is False. + max_spots_per_chunk : int, optional + Retained for backwards compatibility and ignored; transcript + aggregation is now streamed. See `create_pseudo_spots`. Returns ------- @@ -343,6 +567,10 @@ def xenium_to_spatialdata( str(dataset_path), aligned_images=False, morphology_focus=False, n_jobs=n_jobs, cells_as_circles=False ) + # Older "Preview"/"With_Addon" bundles store cell_id and fov_name as + # parquet binary, which dask cannot hand to pyarrow when writing points. + sdata["transcripts"] = decode_bytes_columns(sdata["transcripts"]) + # Fix any validation errors in the table sdata["table"] = fix_table_validation_errors(sdata["table"]) @@ -362,13 +590,31 @@ def xenium_to_spatialdata( logger.info("Adding aligned H&E image...") try: sdata = sd.read_zarr(zarr_path) + + # Parse the alignment once, here, and apply it to both the image + # and the tissue contours below. `xenium_aligned_image` would + # otherwise build its own transform with a raw `pd.read_csv`, + # which mis-reads keypoint-format alignment files -- leaving + # `he_image` and `tissue_contours` on two different transforms. + if file_paths["alignment_file_path"] is not None: + alignment = read_xenium_alignment(str(file_paths["alignment_file_path"])) + transform = Affine(alignment, input_axes=("x", "y"), output_axes=("x", "y")) + else: + logger.warning( + "No alignment file found next to %s; H&E image and tissue contours " + "will use an identity transform and will not be aligned to the " + "Xenium coordinate space.", + file_paths["img_path"].name, + ) + transform = Identity() + aligned_he_image = xenium_aligned_image( image_path=str(file_paths["img_path"]), - alignment_file=str(file_paths["alignment_file_path"]) - if file_paths["alignment_file_path"] - else None, + alignment_file=None, # applied below, from read_xenium_alignment + dims=_aligned_image_dims(str(file_paths["img_path"])), image_models_kwargs={"scale_factors": [2, 2], "chunks": {"x": 512, "y": 512}}, ) + set_transformation(aligned_he_image, transform, to_coordinate_system="global") sdata["he_image"] = aligned_he_image sdata.write_element("he_image") @@ -384,17 +630,9 @@ def xenium_to_spatialdata( # at a different pixel size). Pre-scaling by it here would double-apply # a scale factor and place `tissue_contours` far from the true tissue. tissue_contours = segment_tissue(wsi_path=str(file_paths["img_path"]), pixel_size=1.0, method="otsu") - if file_paths["alignment_file_path"] is not None: - # `read_xenium_alignment` (not a raw `pd.read_csv`) so that Xenium - # Explorer >= v2.0 keypoint-format alignment files (columns - # `fixedX`/`fixedY`/`alignmentX`/`alignmentY`) are converted to an - # affine correctly instead of being fed to `Affine` verbatim. - alignment = read_xenium_alignment(str(file_paths["alignment_file_path"])) - transform = Affine(alignment, input_axes=("x", "y"), output_axes=("x", "y")) - else: - logger.warning("No alignment file found; tissue contours will use an identity transform.") - transform = Identity() + # Same `transform` the image got, so the contours cannot drift + # from the pixels they were segmented from. sdata["tissue_contours"] = ShapesModel.parse(tissue_contours, transformations={"global": transform}) sdata.write_element("tissue_contours") sdata.write_metadata() @@ -412,7 +650,7 @@ def xenium_to_spatialdata( for spot_size in spot_sizes: logger.info("Processing spot size: %sµm", spot_size) - sdata_iter = create_pseudo_spots(sdata_iter, spot_size, overlap, values) + sdata_iter = create_pseudo_spots(sdata_iter, spot_size, overlap, values, max_spots_per_chunk) logger.info("Successfully created %s", zarr_path) @@ -431,6 +669,7 @@ def xenium_to_spatialdata_zip( n_jobs: int = 1, overwrite: bool = False, keep_zarr: bool = True, + max_spots_per_chunk: int | None = 50_000, ) -> Path: """ Convert Xenium data to SpatialData zarr and create a zip archive. @@ -467,6 +706,9 @@ def xenium_to_spatialdata_zip( Whether to overwrite existing files. Default is False. keep_zarr : bool, optional Whether to keep the unzipped zarr directory after creating zip. Default is True. + max_spots_per_chunk : int, optional + Retained for backwards compatibility and ignored; transcript + aggregation is now streamed. See `create_pseudo_spots`. Returns ------- @@ -501,6 +743,7 @@ def xenium_to_spatialdata_zip( values=values, n_jobs=n_jobs, overwrite=overwrite, + max_spots_per_chunk=max_spots_per_chunk, ) # Create zip archive @@ -602,6 +845,7 @@ def download_xenium_study( timeout: int = 180, extract: bool = True, overwrite: bool = False, + resume: bool = True, dry_run: bool = False, ) -> list[DownloadResult]: """Download a Xenium study's raw asset bundle. @@ -628,6 +872,10 @@ def download_xenium_study( (`*_xe_outs.zip` is always skipped). Default True. overwrite : bool, optional Re-download assets that already exist on disk. Default False. + resume : bool, optional + Continue a partially-downloaded asset from its leftover `.part` + file using an HTTP `Range` request instead of refetching it from + the beginning. Default True. dry_run : bool, optional Resolve and report what would be downloaded without any network activity. Default False. @@ -648,6 +896,7 @@ def download_xenium_study( timeout=timeout, extract=extract, overwrite=overwrite, + resume=resume, dry_run=dry_run, ) return downloader.run(source, studies=studies, kinds=kinds) diff --git a/tests/test_utils.py b/tests/test_utils.py index 348a950..f0e6479 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -15,8 +15,12 @@ _mask_to_gdf as mask_to_gdf, ) from spatialrefinery.core.utils import ( + assign_points_to_hexes, bin_centroids, + bin_points_to_hex_counts, + hex_candidate_offsets, hex_grid_centroids, + hex_lattice_params, human_bytes, parse_curl_manifest, safe_extract_zip, @@ -163,6 +167,166 @@ def test_hex_grid_centroids_invalid_overlap_raises(overlap: float) -> None: hex_grid_centroids((0.0, 0.0, 100.0, 100.0), spot_size_um=20.0, overlap=overlap) +# --------------------------------------------------------------------- # +# Analytic hex membership (replaces the shapely sjoin in the transcript +# aggregation -- these are the tests that pin it to the old semantics) +# --------------------------------------------------------------------- # + + +def _shapely_reference(x, y, bounds, spot_size_um, overlap): + """Ground truth: build the polygons and sjoin, exactly as before.""" + geopandas = pytest.importorskip("geopandas") + shapely = pytest.importorskip("shapely.geometry") + + from spatialrefinery.core.utils import hexagon_vertices + + centroids = hex_grid_centroids(bounds, spot_size_um, overlap) + hexes = geopandas.GeoDataFrame( + geometry=[shapely.Polygon(hexagon_vertices(cx, cy, spot_size_um / 2.0)) for cx, cy in centroids] + ) + points = geopandas.GeoDataFrame(geometry=geopandas.points_from_xy(x, y)) + joined = hexes.sjoin(points) + return set(zip(joined.index_right.tolist(), joined.index.tolist(), strict=True)) + + +@pytest.mark.parametrize( + ("spot_size_um", "overlap"), + [(55.0, 0.0), (55.0, 0.06), (100.0, 0.06), (55.0, 0.25), (30.0, 0.4)], +) +def test_assign_points_to_hexes_matches_shapely_sjoin(spot_size_um: float, overlap: float) -> None: + """The analytic test must reproduce `sjoin` pair-for-pair. + + This is the contract that lets `_aggregate_transcripts_hex` stand in for + `SpatialData.aggregate`: same pairs means same counts, including the + duplicates that overlapping hexagons produce. + """ + bounds = (0.0, 0.0, 600.0, 600.0) + rng = np.random.default_rng(11) + x = rng.uniform(-spot_size_um, 600.0 + spot_size_um, 20_000) + y = rng.uniform(-spot_size_um, 600.0 + spot_size_um, 20_000) + + lattice = hex_lattice_params(bounds, spot_size_um, overlap) + positions, spot_ids = assign_points_to_hexes(x, y, lattice, overlap) + + assert set(zip(positions.tolist(), spot_ids.tolist(), strict=True)) == _shapely_reference( + x, y, bounds, spot_size_um, overlap + ) + + +def test_assign_points_to_hexes_covers_the_plane_without_overlap() -> None: + """With overlap=0 the hexagons tile, so every interior point lands in exactly one.""" + bounds = (0.0, 0.0, 400.0, 400.0) + lattice = hex_lattice_params(bounds, 55.0, 0.0) + rng = np.random.default_rng(3) + x = rng.uniform(50.0, 350.0, 5_000) + y = rng.uniform(50.0, 350.0, 5_000) + + positions, _ = assign_points_to_hexes(x, y, lattice, 0.0) + assert np.array_equal(np.bincount(positions, minlength=5_000), np.ones(5_000, dtype=int)) + + +def test_assign_points_to_hexes_overlap_assigns_some_points_twice() -> None: + """Overlapping spots genuinely double-count, which the sjoin also did.""" + bounds = (0.0, 0.0, 400.0, 400.0) + lattice = hex_lattice_params(bounds, 55.0, 0.06) + rng = np.random.default_rng(5) + x = rng.uniform(50.0, 350.0, 5_000) + y = rng.uniform(50.0, 350.0, 5_000) + + positions, _ = assign_points_to_hexes(x, y, lattice, 0.06) + multiplicity = np.bincount(positions, minlength=5_000) + assert multiplicity.min() >= 1 + assert multiplicity.max() > 1 + + +def test_hex_grid_centroids_matches_lattice_params_indexing() -> None: + """Spot id `i * nx + j` must address the centroid the lattice describes.""" + bounds = (-13.0, 7.5, 500.0, 300.0) + lattice = hex_lattice_params(bounds, 55.0, 0.06) + centroids = hex_grid_centroids(bounds, 55.0, 0.06) + + assert len(centroids) == lattice["nx"] * lattice["ny"] + for row in (0, 1, lattice["ny"] - 1): + for col in (0, lattice["nx"] - 1): + expected_x = lattice["x_range"][col] + (lattice["dx"] / 2 if row % 2 else 0.0) + assert centroids[row * lattice["nx"] + col] == pytest.approx((expected_x, lattice["y_range"][row])) + + +@pytest.mark.parametrize( + ("overlap", "rows", "cols"), + [(0.0, [0, 1], [0, 1]), (0.06, [0, 1], [0, 1]), (0.4, [-1, 0, 1, 2], [0, 1])], +) +def test_hex_candidate_offsets_widen_with_overlap(overlap, rows, cols) -> None: + row_offsets, col_offsets = hex_candidate_offsets(overlap) + assert list(row_offsets) == rows + assert list(col_offsets) == cols + + +def test_bin_points_to_hex_counts_matches_a_dense_reference() -> None: + """Streaming in batches must give the same matrix as counting in one go.""" + bounds = (0.0, 0.0, 300.0, 300.0) + overlap, spot_size, n_genes = 0.06, 55.0, 7 + lattice = hex_lattice_params(bounds, spot_size, overlap) + n_spots = lattice["nx"] * lattice["ny"] + + rng = np.random.default_rng(17) + x = rng.uniform(0.0, 300.0, 30_000) + y = rng.uniform(0.0, 300.0, 30_000) + codes = rng.integers(0, n_genes, 30_000) + + counts = bin_points_to_hex_counts( + ((x[s], y[s], codes[s]) for s in (slice(0, 7_000), slice(7_000, 21_000), slice(21_000, None))), + lattice, + n_spots=n_spots, + n_genes=n_genes, + overlap=overlap, + ) + + expected = np.zeros((n_spots, n_genes), dtype=np.int64) + positions, spot_ids = assign_points_to_hexes(x, y, lattice, overlap) + np.add.at(expected, (spot_ids, codes[positions]), 1) + + assert np.array_equal(counts.toarray(), expected) + assert counts.sum() == len(positions) + + +def test_bin_points_to_hex_counts_flushes_without_changing_the_result() -> None: + """`flush_pairs` bounds the merge buffer only; it must not alter counts.""" + bounds = (0.0, 0.0, 200.0, 200.0) + lattice = hex_lattice_params(bounds, 55.0, 0.06) + n_spots = lattice["nx"] * lattice["ny"] + + rng = np.random.default_rng(23) + batches = [(rng.uniform(0, 200, 4_000), rng.uniform(0, 200, 4_000), rng.integers(0, 5, 4_000)) for _ in range(5)] + + unflushed = bin_points_to_hex_counts(iter(batches), lattice, n_spots, 5, 0.06, flush_pairs=10**9) + flushed = bin_points_to_hex_counts(iter(batches), lattice, n_spots, 5, 0.06, flush_pairs=100) + assert (unflushed != flushed).nnz == 0 + + +def test_bin_points_to_hex_counts_drops_negative_gene_codes() -> None: + """A -1 code is how an unmapped feature is skipped, not counted into gene 0.""" + bounds = (0.0, 0.0, 200.0, 200.0) + lattice = hex_lattice_params(bounds, 55.0, 0.0) + n_spots = lattice["nx"] * lattice["ny"] + + x = np.array([100.0, 100.0, 100.0]) + y = np.array([100.0, 100.0, 100.0]) + kept = bin_points_to_hex_counts([(x, y, np.array([0, 0, 0]))], lattice, n_spots, 3, 0.0) + dropped = bin_points_to_hex_counts([(x, y, np.array([0, -1, -1]))], lattice, n_spots, 3, 0.0) + + assert kept.sum() == 3 + assert dropped.sum() == 1 + + +def test_bin_points_to_hex_counts_empty_input_returns_empty_matrix() -> None: + lattice = hex_lattice_params((0.0, 0.0, 200.0, 200.0), 55.0, 0.06) + n_spots = lattice["nx"] * lattice["ny"] + counts = bin_points_to_hex_counts([], lattice, n_spots, 4, 0.06) + assert counts.shape == (n_spots, 4) + assert counts.nnz == 0 + + # --------------------------------------------------------------------- # # Tissue mask -> polygons (regression net for the hestcore -> cv2 rewrite) # --------------------------------------------------------------------- # @@ -219,3 +383,94 @@ def test_fix_table_validation_errors_renames_invalid_var_columns(adata) -> None: assert "bad col" not in fixed.var.columns assert "bad_col" in fixed.var.columns validate_table_attr_keys(fixed) # must no longer raise + + +# --------------------------------------------------------------------- # +# Points: bytes-column decoding +# --------------------------------------------------------------------- # + + +def _bytes_points(): + """A points frame shaped like an older Xenium `transcripts.parquet`. + + `cell_id` / `fov_name` hold `bytes` under a bare `object` dtype, which is + what the "Preview" / "With_Addon" bundles produce. + + `convert-string` is switched off while the frame is built because + `dd.from_pandas` otherwise re-types object columns to `StringDtype` and the + bug disappears. The real frames come from `dd.read_parquet` with a declared + object-dtype meta, which is what this reproduces. + """ + import dask + import dask.dataframe as dd + import pandas as pd + from spatialdata.models import PointsModel + + df = pd.DataFrame( + { + "x": np.array([1.0, 2.0], dtype="float32"), + "y": np.array([3.0, 4.0], dtype="float32"), + "cell_id": np.array([b"UNASSIGNED", b"abcd-1"], dtype=object), + "fov_name": np.array([b"A5", b"B7"], dtype=object), + } + ) + with dask.config.set({"dataframe.convert-string": False}): + return PointsModel.parse(dd.from_pandas(df, npartitions=1)) + + +def test_decode_bytes_columns_makes_points_parquet_writable() -> None: + """The decoded frame's dummy meta must be inferable by pyarrow. + + This is the actual failure mode: `to_parquet` types the schema from + `meta_nonempty`, which fills a bare `object` column with an `object()` + sentinel that pyarrow rejects. + """ + import pyarrow as pa + from dask.dataframe.utils import meta_nonempty + + from spatialrefinery.core.utils import decode_bytes_columns + + points = _bytes_points() + + with pytest.raises(pa.ArrowInvalid): + pa.Schema.from_pandas(meta_nonempty(points._meta), preserve_index=False) + + decoded = decode_bytes_columns(points) + + pa.Schema.from_pandas(meta_nonempty(decoded._meta), preserve_index=False) + + +def test_decode_bytes_columns_decodes_values_and_keeps_transformations() -> None: + """Values become real strings and the coordinate transforms survive.""" + from spatialdata.transformations import Scale, get_transformation, set_transformation + + from spatialrefinery.core.utils import decode_bytes_columns + + points = _bytes_points() + set_transformation(points, Scale([2.0, 2.0], axes=("x", "y")), to_coordinate_system="global") + before = get_transformation(points, get_all=True) + + decoded = decode_bytes_columns(points) + result = decoded.compute() + + assert list(result["cell_id"]) == ["UNASSIGNED", "abcd-1"] + assert list(result["fov_name"]) == ["A5", "B7"] + assert str(get_transformation(decoded, get_all=True)) == str(before) + + +def test_decode_bytes_columns_no_object_columns_is_a_noop() -> None: + """A frame with no `object` columns is returned untouched.""" + import dask.dataframe as dd + import pandas as pd + from spatialdata.models import PointsModel + + from spatialrefinery.core.utils import decode_bytes_columns + + points = PointsModel.parse( + dd.from_pandas( + pd.DataFrame({"x": np.array([1.0], dtype="float32"), "y": np.array([2.0], dtype="float32")}), + npartitions=1, + ) + ) + + assert decode_bytes_columns(points) is points diff --git a/tests/test_xenium_io.py b/tests/test_xenium_io.py new file mode 100644 index 0000000..1cee79c --- /dev/null +++ b/tests/test_xenium_io.py @@ -0,0 +1,141 @@ +"""Tests for `spatialrefinery.io.xenium`'s reader-facing helpers. + +Covers the H&E channel-layout probe, which exists because +`spatialdata_io.xenium_aligned_image` infers the axis order from the array +shape alone and asserts on anything that is not 10x's interleaved layout. +Slides are synthesised as small OME-TIFFs written in each layout. +""" + +from __future__ import annotations + +from pathlib import Path + +import numpy as np +import pytest +import tifffile as tf + +from spatialrefinery.io.xenium import _aligned_image_dims + + +def _write_rgb(path: Path, planar: bool) -> Path: + """Write a small RGB OME-TIFF in either planar or interleaved layout.""" + image = np.random.default_rng(0).integers(0, 255, (3, 64, 96), dtype=np.uint8) + if planar: + tf.imwrite(path, image, photometric="rgb", planarconfig="separate", metadata={"axes": "SYX"}) + else: + tf.imwrite( + path, + np.moveaxis(image, 0, -1), + photometric="rgb", + planarconfig="contig", + metadata={"axes": "YXS"}, + ) + return path + + +def test_aligned_image_dims_planar_returns_explicit_dims(tmp_path: Path) -> None: + """`planarconfig=SEPARATE` reads back as (1, c, y, x), which the reader mis-parses.""" + path = _write_rgb(tmp_path / "planar.ome.tif", planar=True) + + from dask_image.imread import imread + + assert imread(str(path)).shape == (1, 3, 64, 96) + assert _aligned_image_dims(str(path)) == ("dummy", "c", "y", "x") + + +def test_aligned_image_dims_interleaved_defers_to_the_reader(tmp_path: Path) -> None: + """10x's own layout is (1, y, x, c); the reader already handles it, so return None.""" + path = _write_rgb(tmp_path / "interleaved.ome.tif", planar=False) + + from dask_image.imread import imread + + assert imread(str(path)).shape == (1, 64, 96, 3) + assert _aligned_image_dims(str(path)) is None + + +def test_aligned_image_dims_rejects_an_unrecognisable_layout(tmp_path: Path) -> None: + """A single-channel image has no channel axis to find; fail loudly, not with an assert.""" + path = tmp_path / "grey.ome.tif" + tf.imwrite(path, np.zeros((64, 96), dtype=np.uint8)) + + with pytest.raises(ValueError, match="Cannot infer the channel axis"): + _aligned_image_dims(str(path)) + + +def _tiny_sdata(spot_size: float = 55.0, overlap: float = 0.06): + """A minimal SpatialData with transcripts and a matching hex-spot element.""" + import dask.dataframe as dd + import pandas as pd + import spatialdata as sd + from spatialdata.models import PointsModel, ShapesModel + + from spatialrefinery.core.utils import create_hexagonal_spots, hex_lattice_params, xy_bounds + + rng = np.random.default_rng(4) + n = 5_000 + genes = ["GeneA", "GeneB", "GeneC"] + frame = pd.DataFrame( + { + "x": rng.uniform(0.0, 400.0, n), + "y": rng.uniform(0.0, 400.0, n), + "feature_name": pd.Categorical(rng.choice(genes, n), categories=genes), + } + ) + points = PointsModel.parse(dd.from_pandas(frame, npartitions=3)) + + bounds = xy_bounds(points, "x", "y") + lattice = hex_lattice_params(bounds, spot_size, overlap) + spots = ShapesModel.parse(create_hexagonal_spots(points, spot_size_um=spot_size, overlap=overlap, bounds=bounds)) + sdata = sd.SpatialData(points={"transcripts": points}, shapes={"spots_55um": spots}) + return sdata, lattice, overlap + + +def test_aggregate_transcripts_hex_returns_a_table_annotating_the_spots() -> None: + """The table must be linked to its spots element, not just carry the right counts. + + Regression test: returning a bare `AnnData` left the counts correct but + unattached, so consumers reported the spots element as having no annotating + tables. `SpatialData.aggregate`, which this replaced, returned a parsed one. + """ + from spatialrefinery.io.xenium import _aggregate_transcripts_hex + + sdata, lattice, overlap = _tiny_sdata() + table = _aggregate_transcripts_hex(sdata, "spots_55um", lattice, overlap) + + assert table.uns["spatialdata_attrs"] == { + "region": "spots_55um", + "region_key": "region", + "instance_key": "instance_id", + } + assert list(table.obs["region"].cat.categories) == ["spots_55um"] + assert np.array_equal(table.obs["instance_id"].to_numpy(), sdata["spots_55um"].index.to_numpy()) + assert table.obs.index.equals(sdata["spots_55um"].index.astype(str)) + + # and it must survive being attached to the SpatialData object + sdata["spots_55um_table"] = table + assert "spots_55um_table" in sdata.tables + + +def test_aggregate_transcripts_hex_counts_every_transcript() -> None: + """Counts must be conserved: each transcript lands in >=1 spot, overlaps included.""" + from spatialrefinery.core.utils import assign_points_to_hexes + from spatialrefinery.io.xenium import _aggregate_transcripts_hex + + sdata, lattice, overlap = _tiny_sdata() + table = _aggregate_transcripts_hex(sdata, "spots_55um", lattice, overlap) + + frame = sdata["transcripts"][["x", "y"]].compute() + positions, _ = assign_points_to_hexes(frame["x"].to_numpy(), frame["y"].to_numpy(), lattice, overlap) + assert table.X.sum() == len(positions) + + +def test_aggregate_transcripts_hex_raises_when_nothing_is_counted() -> None: + """An all-zero result means a broken mapping; it must not be written silently.""" + from spatialrefinery.core.utils import hex_lattice_params + from spatialrefinery.io.xenium import _aggregate_transcripts_hex + + sdata, _, overlap = _tiny_sdata() + # a lattice somewhere else entirely: no transcript can fall in any spot + elsewhere = hex_lattice_params((1e6, 1e6, 1e6 + 400.0, 1e6 + 400.0), 55.0, overlap) + with pytest.raises(ValueError, match="produced no counts at all"): + _aggregate_transcripts_hex(sdata, "spots_55um", elsewhere, overlap)