From 3ccc6786aaacd17d5dacd4bb68a22cb595c0e9dd Mon Sep 17 00:00:00 2001 From: Rushin Gindra Date: Fri, 11 Sep 2026 18:31:05 +0200 Subject: [PATCH 1/2] Keep the batch dim on batches of one and accept a SpatialData object in SpatialDataset Two output-preserving changes needed to run inference over many stores in blocks: - `FlowPipeline.__call__` squeezed every singleton dim of the sampled `(B, n_genes, 1)` tensor, so a batch of size one (the last block of a store) collapsed to `(n_genes,)` and broke the concatenation and the `* std + mean` broadcast. `squeeze(-1)` drops only the gene-token dim; for B > 1 the result is identical. - `SpatialDataset` accepts an already-read `SpatialData` object as `zarr_path` and uses it as is, so a caller can `sd.read_zarr(path, selection=("images", "shapes", "tables"))` once (skipping the large `transcripts` element) and share it with the dataset, the ground truth and the write-back. Paths still go through `sd.read_zarr(zarr_path)`. Tests: `tests/conftest.py` gains `make_synthetic_store`, which writes a minimal SpatialData zarr (multiscale `he_image`, circular `nucleus_boundaries`, Poisson-count `table`, identity transforms) and a `synthetic_store` fixture. The placeholder `SpatialDataset` test is replaced by real tests of length, the demo image transform, the blank patch for a border cell, construction from a `SpatialData` object with `selection=`, and gene-list subsetting; `test_inference.py` adds `test_flow_pipeline_handles_batch_size_one`, which fails against the previous `squeeze()`. Verified on CPU against 2436b23: `pytest tests` 33 passed, 4 skipped (the apex/flash-attn-only modules, as before); `ruff check` and `ruff format --check` clean on `src` and `tests`. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_016zPfR5j2rpSaamQUrLmRuJ --- src/phoenix/datasets/zarr_dataset.py | 11 ++-- src/phoenix/helpers/inference.py | 2 +- tests/conftest.py | 75 ++++++++++++++++++++++++++++ tests/test_datasets.py | 67 +++++++++++++++++++++---- tests/test_inference.py | 20 +++++++- 5 files changed, 158 insertions(+), 17 deletions(-) diff --git a/src/phoenix/datasets/zarr_dataset.py b/src/phoenix/datasets/zarr_dataset.py index 57232a4..bcbb933 100644 --- a/src/phoenix/datasets/zarr_dataset.py +++ b/src/phoenix/datasets/zarr_dataset.py @@ -4,6 +4,8 @@ © Peng Lab / Helmholtz Munich """ +from pathlib import Path + import numpy as np import spatialdata as sd import torch @@ -28,7 +30,8 @@ class SpatialDataset(Dataset): Parameters ---------- zarr_path - Path to a SpatialData ``.zarr`` store. + Path to a SpatialData ``.zarr`` store, or an already-read ``SpatialData`` + object (e.g. from ``sd.read_zarr(path, selection=...)``) used as is. table_type Key of the anndata table to read from the store (used when `adata_transform` is not given). @@ -48,7 +51,7 @@ class SpatialDataset(Dataset): def __init__( self, - zarr_path: str, + zarr_path: str | Path | sd.SpatialData, table_type: str, gene_list: list, patch_size: int = 224, @@ -56,8 +59,8 @@ def __init__( adata_transform: Compose | None = None, image_transform: Compose | None = None, ): - # read zarr file with spatialdata - self.sdata = sd.read_zarr(zarr_path) + # read zarr file with spatialdata, unless an already-read store is given + self.sdata = zarr_path if isinstance(zarr_path, sd.SpatialData) else sd.read_zarr(zarr_path) if adata_transform: adata = self.sdata[table_type] diff --git a/src/phoenix/helpers/inference.py b/src/phoenix/helpers/inference.py index 2cd1253..dc2afb6 100644 --- a/src/phoenix/helpers/inference.py +++ b/src/phoenix/helpers/inference.py @@ -151,7 +151,7 @@ def __call__(self, gene_list: list, dataloader: DataLoader): device=device, ) - gex_pred = gex_pred.float().squeeze().detach().cpu().numpy() + gex_pred = gex_pred.float().squeeze(-1).detach().cpu().numpy() pred_list.append(gex_pred) coords_list.append(coords) diff --git a/tests/conftest.py b/tests/conftest.py index 4cee7b4..3f790d3 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,6 +1,11 @@ +from collections.abc import Sequence +from pathlib import Path + import anndata as ad import numpy as np +import pandas as pd import pytest +from scipy.sparse import csr_matrix @pytest.fixture @@ -10,3 +15,73 @@ def adata(): adata.var_names = ["PECAM1", "MMRN2"] adata.obsm["spatial"] = np.array([[0, 0], [10, 10], [20, 20]], dtype=np.float32) return adata + + +def make_synthetic_store( + path: Path, + n_cells: int = 32, + genes: Sequence[str] = ("PECAM1", "MMRN2", "MYH11", "SFRP2"), + image_size: int = 512, + seed: int = 0, +) -> Path: + """ + Write a minimal SpatialData ``.zarr`` store that `SpatialDataset` can consume. + + The store holds a multiscale ``he_image``, circular ``nucleus_boundaries`` and a + ``table`` of Poisson counts annotating them, all under identity transforms, so the + native resolution resolves to 1.0 micron per pixel. Cells are placed away from the + image border except the last one, which sits on the left edge so its patch is clipped + along x only and comes out non-square (a cell in the corner would clip both axes to an + empty, square patch that escapes the dataset's blank-patch fallback). + + Parameters + ---------- + path + Where to write the store. + n_cells + Number of cells (rows of the table, nucleus shapes). + genes + Gene panel of the table, in this order. + image_size + Side length, in pixels, of the square ``he_image`` at scale 0. + seed + Seed of the ``numpy`` generator drawing coordinates, pixels and counts. + + Returns + ------- + The path the store was written to. + """ + sd = pytest.importorskip("spatialdata") + from spatialdata.models import Image2DModel, ShapesModel, TableModel + + rng = np.random.default_rng(seed) + genes = list(genes) + + xy = rng.uniform(64, image_size - 64, size=(n_cells, 2)) + xy[-1] = (2.0, image_size / 2) + + image = rng.integers(0, 256, size=(3, image_size, image_size), dtype=np.uint8) + he_image = Image2DModel.parse(image, dims=("c", "y", "x"), scale_factors=[2]) + nuclei = ShapesModel.parse(xy, geometry=0, radius=5.0) + + counts = csr_matrix(rng.poisson(2.0, size=(n_cells, len(genes))).astype(np.float32)) + obs = pd.DataFrame( + { + "instance_id": np.arange(n_cells), + "region": pd.Categorical(["nucleus_boundaries"] * n_cells), + }, + index=np.arange(n_cells).astype(str), + ) + adata = ad.AnnData(X=counts, obs=obs, var=pd.DataFrame(index=genes), obsm={"spatial": xy}) + table = TableModel.parse(adata, region="nucleus_boundaries", region_key="region", instance_key="instance_id") + + sdata = sd.SpatialData( + images={"he_image": he_image}, shapes={"nucleus_boundaries": nuclei}, tables={"table": table} + ) + sdata.write(path) + return path + + +@pytest.fixture +def synthetic_store(tmp_path): + return make_synthetic_store(tmp_path / "store.zarr") diff --git a/tests/test_datasets.py b/tests/test_datasets.py index 8f7b6c0..110d943 100644 --- a/tests/test_datasets.py +++ b/tests/test_datasets.py @@ -4,8 +4,13 @@ h5py = pytest.importorskip("h5py") pytest.importorskip("torch") pytest.importorskip("PIL") +sd = pytest.importorskip("spatialdata") + +import torch # noqa: E402 +from torchvision.transforms import InterpolationMode, v2 # noqa: E402 from phoenix.datasets.h5py_dataset import H5PYDataset # noqa: E402 +from phoenix.datasets.zarr_dataset import SpatialDataset # noqa: E402 @pytest.fixture @@ -44,16 +49,56 @@ def transform(img): assert isinstance(patch, np.ndarray) -def test_spatial_dataset_is_a_torch_dataset(): - # Instantiating SpatialDataset needs a real SpatialData .zarr store (he_image + - # nucleus_boundaries elements, a table with a matching gene panel), which isn't - # practical to fabricate in a unit test; this at least confirms the class is - # importable and conforms to the expected torch Dataset interface. - pytest.importorskip("spatialdata") - from torch.utils.data import Dataset +GENES = ["PECAM1", "MMRN2", "MYH11", "SFRP2"] + +# the demo notebook's image transform +DEMO_TRANSFORM = v2.Compose( + [ + v2.Resize((224, 224), interpolation=InterpolationMode.BICUBIC), + v2.CenterCrop((224, 224)), + v2.ToTensor(), + v2.Normalize((0.707223, 0.578729, 0.703617), (0.211883, 0.230117, 0.177517)), + ] +) + + +def test_spatial_dataset_length(synthetic_store): + dataset = SpatialDataset(synthetic_store, "table", GENES) + assert len(dataset) == dataset.sdata["table"].n_obs == 32 + + +def test_spatial_dataset_getitem_applies_transform(synthetic_store): + dataset = SpatialDataset(synthetic_store, "table", GENES, image_transform=DEMO_TRANSFORM) + image, coords = dataset[0] + assert isinstance(image, torch.Tensor) + assert image.shape == (3, 224, 224) + assert image.dtype == torch.float32 + np.testing.assert_array_equal(coords, dataset.adata.obsm["spatial"][0].astype(int)) + + +def test_spatial_dataset_border_cell_yields_blank_patch(synthetic_store): + dataset = SpatialDataset(synthetic_store, "table", GENES) + image, _ = dataset[len(dataset) - 1] + image = np.asarray(image) + assert image.shape == (224, 224, 3) + assert image.dtype == np.uint8 + assert image.max() == 0 + + +def test_spatial_dataset_accepts_spatialdata_object(synthetic_store): + sdata = sd.read_zarr(synthetic_store, selection=("images", "shapes", "tables")) + from_path = SpatialDataset(synthetic_store, "table", GENES, image_transform=DEMO_TRANSFORM) + from_sdata = SpatialDataset(sdata, "table", GENES, image_transform=DEMO_TRANSFORM) + assert from_sdata.sdata is sdata + assert len(from_sdata) == len(from_path) + + image_path, coords_path = from_path[0] + image_sdata, coords_sdata = from_sdata[0] + np.testing.assert_array_equal(coords_sdata, coords_path) + assert torch.equal(image_sdata, image_path) - from phoenix.datasets.zarr_dataset import SpatialDataset - assert issubclass(SpatialDataset, Dataset) - assert hasattr(SpatialDataset, "__getitem__") - assert hasattr(SpatialDataset, "__len__") +def test_spatial_dataset_subsets_to_gene_list(synthetic_store): + dataset = SpatialDataset(synthetic_store, "table", ["MYH11", "PECAM1"]) + assert dataset.adata.var_names.tolist() == ["MYH11", "PECAM1"] + assert dataset.gene_matrix.shape == (len(dataset), 2) diff --git a/tests/test_inference.py b/tests/test_inference.py index d50f416..39d7f78 100644 --- a/tests/test_inference.py +++ b/tests/test_inference.py @@ -75,7 +75,6 @@ def test_flow_pipeline_runs_on_cpu(pipeline_model): n_samples, n_genes = 4, 3 feats = torch.randn(n_samples, 5, 16) coords = torch.zeros(n_samples, 2) - # batch_size=1 would let the squeeze() inside __call__ collapse the batch dim loader = DataLoader(TensorDataset(feats, coords), batch_size=2) stats = {"mean": np.zeros(n_genes), "std": np.ones(n_genes)} @@ -89,6 +88,25 @@ def test_flow_pipeline_runs_on_cpu(pipeline_model): assert sum(len(c) for c in coords_list) == n_samples +def test_flow_pipeline_handles_batch_size_one(pipeline_model): + """A batch of one must keep its batch dim: `__call__` squeezes only the trailing gene-token dim.""" + from torch.utils.data import DataLoader, TensorDataset + + torch.manual_seed(0) + n_samples, n_genes = 3, 2 + feats = torch.randn(n_samples, 5, 16) + coords = torch.zeros(n_samples, 2) + loader = DataLoader(TensorDataset(feats, coords), batch_size=1) + + stats = {"mean": np.zeros(n_genes), "std": np.ones(n_genes)} + pipeline = FlowPipeline(model=pipeline_model, stats=stats, atol=1e-1, rtol=1e-1) + + gex_pred, coords_list = pipeline(["A", "B"], loader) + + assert gex_pred.shape == (n_samples, n_genes) + assert len(coords_list) == n_samples + + def test_run_fast_flow_matches_run_flow(tiny_model): """The K/V-caching sampler must produce the same output as run_flow.""" from phoenix.helpers.fast_sampler import run_fast_flow From 6950174af633cc374720d7a103a4a500ac973066 Mon Sep 17 00:00:00 2001 From: Rushin Gindra Date: Sun, 13 Sep 2026 16:54:45 +0200 Subject: [PATCH 2/2] Set the synthetic store's obsm after construction to satisfy mypy MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The pre-commit CI job failed on its mypy hook: `AnnData.__init__` annotates `obsm` as `Mapping[str, Sequence[Any]]`, and the `(n_cells, 2)` coordinate array passed there is not a `Sequence`, so `hatch check types` reported `tests/conftest.py:75: error: Dict entry 0 has incompatible type` and the job — and with it the "Tests pass in all hatch environments" gate — went red. Assigning through `adata.obsm["spatial"]` instead goes through `AxisArrays`, which is typed for arrays, and matches how the existing `adata` fixture in this file sets the same key. The written store is unchanged. Verified on CPU in the hatch envs: `hatch check types` clean on 22 files, `prek run --all-files` (the exact CI invocation) all 13 hooks pass, `hatch test` 33 passed, 4 skipped. Co-Authored-By: Claude Opus 5 (1M context) --- tests/conftest.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/conftest.py b/tests/conftest.py index 3f790d3..1f5b223 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -72,7 +72,8 @@ def make_synthetic_store( }, index=np.arange(n_cells).astype(str), ) - adata = ad.AnnData(X=counts, obs=obs, var=pd.DataFrame(index=genes), obsm={"spatial": xy}) + adata = ad.AnnData(X=counts, obs=obs, var=pd.DataFrame(index=genes)) + adata.obsm["spatial"] = xy table = TableModel.parse(adata, region="nucleus_boundaries", region_key="region", instance_key="instance_id") sdata = sd.SpatialData(