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..1f5b223 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,74 @@ 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)) + adata.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