Keep the batch dim on batches of one and accept a SpatialData object in SpatialDataset

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,16 +51,16 @@ 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,
         target_mpp: float = 0.5,
         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]

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)
 

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")

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)

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