Fix correctness, robustness and memory issues in the MERSCOPE reader

src/spatialdata_io/readers/merscope.py

@@ -12,14 +12,16 @@
 import geopandas
 import numpy as np
 import pandas as pd
+import shapely
 import xarray
 from dask import array as da
 from dask_image.imread import imread
+from scipy.sparse import csr_matrix
 from shapely.geometry import MultiPolygon
 from spatialdata import SpatialData
 from spatialdata._logging import logger
 from spatialdata.models import Image2DModel, PointsModel, ShapesModel, TableModel
-from spatialdata.transformations import Affine, BaseTransformation
+from spatialdata.transformations import Affine, BaseTransformation, Identity
 
 from spatialdata_io._constants._constants import MerscopeKeys
 from spatialdata_io._docs import inject_docs
@@ -34,10 +36,10 @@
 def _get_channel_names(images_dir: Path) -> list[str]:
     exp = r"mosaic_(?P<stain>[\w|-]+[0-9]?)_z(?P<z>[0-9]+).tif"
     matches = [re.search(exp, file.name) for file in images_dir.iterdir()]
-
     stainings = {match.group("stain") for match in matches if match}
-
-    return list(stainings)
+    # sort for a deterministic channel order: set iteration depends on PYTHONHASHSEED, which would
+    # otherwise make the image channel axis (and its `c_coords`) order change between runs.
+    return sorted(stainings)
 
 
 def _get_file_paths(path: Path, vpt_outputs: Path | str | dict[str, Any] | None) -> tuple[Path, Path, Path]:
@@ -73,9 +75,9 @@ def _get_file_paths(path: Path, vpt_outputs: Path | str | dict[str, Any] | None)
 
     if isinstance(vpt_outputs, dict):
         return (
-            vpt_outputs[MerscopeKeys.VPT_NAME_COUNTS],
-            vpt_outputs[MerscopeKeys.VPT_NAME_OBS],
-            vpt_outputs[MerscopeKeys.VPT_NAME_BOUNDARIES],
+            Path(vpt_outputs[MerscopeKeys.VPT_NAME_COUNTS]),
+            Path(vpt_outputs[MerscopeKeys.VPT_NAME_OBS]),
+            Path(vpt_outputs[MerscopeKeys.VPT_NAME_BOUNDARIES]),
         )
 
     raise ValueError(
@@ -170,10 +172,17 @@ def merscope(
     count_path, obs_path, boundaries_path = _get_file_paths(path, vpt_outputs)
     images_dir = path / MerscopeKeys.IMAGES_DIR
 
-    microns_to_pixels = Affine(
-        np.genfromtxt(images_dir / MerscopeKeys.TRANSFORMATION_FILE), input_axes=("x", "y"), output_axes=("x", "y")
-    )
-    transformations = {"global": microns_to_pixels}
+    # The micron-to-mosaic-pixel transform only ships with the mosaic images. Without it the data stays
+    # in micron space (identity), so partial reads (e.g. table/transcripts only) still work.
+    transform_path = images_dir / MerscopeKeys.TRANSFORMATION_FILE
+    if transform_path.exists():
+        transform: BaseTransformation = Affine(
+            np.genfromtxt(transform_path), input_axes=("x", "y"), output_axes=("x", "y")
+        )
+    else:
+        logger.warning(f"Transformation file {transform_path} not found. Elements will be in micron space.")
+        transform = Identity()
+    transformations = {"global": transform}
 
     vizgen_region = path.name if region_name is None else region_name
     slide_name = path.parent.name if slide_name is None else slide_name
@@ -184,20 +193,21 @@ def merscope(
     images = {}
 
     if mosaic_images:
-        z_layers = [z_layers] if isinstance(z_layers, int) else z_layers or []
-
-        reader = _get_reader(backend)
-
-        stainings = _get_channel_names(images_dir)
-        if stainings:
-            for z_layer in z_layers:
-                images[f"{dataset_id}_z{z_layer}"] = reader(
-                    images_dir,
-                    stainings,
-                    z_layer,
-                    image_models_kwargs,
-                    **imread_kwargs,
-                )
+        if not images_dir.exists():
+            logger.warning(f"Image directory {images_dir} does not exist. Mosaic images are not loaded.")
+        else:
+            z_layers = [z_layers] if isinstance(z_layers, int) else z_layers or []
+            reader = _get_reader(backend)
+            stainings = _get_channel_names(images_dir)
+            if stainings:
+                for z_layer in z_layers:
+                    images[f"{dataset_id}_z{z_layer}"] = reader(
+                        images_dir,
+                        stainings,
+                        z_layer,
+                        image_models_kwargs,
+                        **imread_kwargs,
+                    )
 
     # Transcripts
     points = {}
@@ -304,22 +314,56 @@ def _dask_image_load_merscope(
 
 def _get_points(transcript_path: Path, transformations: dict[str, BaseTransformation]) -> dd.DataFrame:
     transcript_df = dd.read_csv(transcript_path)
+    # Make `gene` a categorical with *known* categories before parsing. Passing an object-dtype column
+    # to `PointsModel.parse` leaves the dask categories "unknown": they then get inferred from a single
+    # partition, which both forces repeated recomputation downstream and can silently turn genes that
+    # are absent from the first partition into NaN. Computing the categories once here avoids both.
+    transcript_df[MerscopeKeys.GENE_KEY] = transcript_df[MerscopeKeys.GENE_KEY].astype("category").cat.as_known()
     transcripts = PointsModel.parse(
         transcript_df,
         coordinates={"x": MerscopeKeys.GLOBAL_X, "y": MerscopeKeys.GLOBAL_Y, "z": MerscopeKeys.GLOBAL_Z},
         transformations=transformations,
         feature_key=MerscopeKeys.GENE_KEY,
+        # `sort=True` skips an eager `is_monotonic_increasing` check on the index that otherwise scans
+        # the whole transcripts file during this nominally lazy read.
+        sort=True,
     )
-    transcripts["gene"] = transcripts["gene"].astype("category")
     return transcripts
 
 
+def _as_multipolygon(geom: Any) -> MultiPolygon:
+    """Valid MultiPolygon of ``geom``'s polygonal parts (empty when it has none).
+
+    Vizgen writes MultiPolygons, but watershed / custom segmentations can contain plain Polygons, and
+    ``make_valid`` can split a self-intersecting ring into a GeometryCollection.
+    """
+    if geom is None or geom.is_empty:
+        return MultiPolygon()
+    if not geom.is_valid:
+        geom = shapely.make_valid(geom)
+    if isinstance(geom, MultiPolygon):
+        return geom
+    parts = geom.geoms if hasattr(geom, "geoms") else [geom]
+    polys = [
+        q
+        for g in parts
+        for q in (g.geoms if isinstance(g, MultiPolygon) else [g])
+        if g.geom_type in ("Polygon", "MultiPolygon")
+    ]
+    return MultiPolygon(polys)
+
+
 def _get_polygons(boundaries_path: Path, transformations: dict[str, BaseTransformation]) -> geopandas.GeoDataFrame:
     geo_df = geopandas.read_parquet(boundaries_path)
     geo_df = geo_df.rename_geometry("geometry")
     geo_df = geo_df[geo_df[MerscopeKeys.Z_INDEX] == 0]  # Avoid duplicate boundaries on all z-levels
-    geo_df = geo_df[geo_df.geometry.is_valid]  # Remove invalid geometries
-    geo_df.geometry = geo_df.geometry.map(lambda x: MultiPolygon(x.geoms))
+    # Repair invalid (e.g. self-intersecting) boundaries instead of dropping them: a dropped polygon leaves
+    # its cell in the table annotating a shape that does not exist.
+    geo_df.geometry = geo_df.geometry.map(_as_multipolygon)
+    empty = geo_df.geometry.is_empty
+    if empty.any():
+        logger.warning(f"Dropping {int(empty.sum())} cell boundaries with no polygonal area after repair.")
+        geo_df = geo_df[~empty]
     geo_df.index = geo_df[MerscopeKeys.METADATA_CELL_KEY].astype(str)
 
     return ShapesModel.parse(geo_df, transformations=transformations)
@@ -337,7 +381,11 @@ def _get_table(
     obs = pd.read_csv(obs_path, index_col=0, dtype={MerscopeKeys.METADATA_CELL_KEY: str})
 
     is_gene = ~data.columns.str.lower().str.contains("blank")
-    adata = anndata.AnnData(data.loc[:, is_gene], obs=obs)
+    genes = data.loc[:, is_gene]
+    adata = anndata.AnnData(genes, obs=obs)
+    # Store counts as a sparse matrix, like the other scverse readers (xenium/cosmx/...): MERSCOPE
+    # matrices are only ~10% dense, so a dense `int64` X wastes roughly an order of magnitude of memory.
+    adata.X = csr_matrix(genes.to_numpy())
 
     assert isinstance(adata.obs, pd.DataFrame)
     adata.obsm["blank"] = data.loc[:, ~is_gene]  # blank fields are excluded from adata.X

tests/test_merscope.py

@@ -1,14 +1,29 @@
+from __future__ import annotations
+
 import sys
-from pathlib import Path
+import warnings
+from typing import TYPE_CHECKING
 
+import numpy as np
+import pandas as pd
 import pytest
+from scipy.sparse import issparse
+from shapely.geometry import MultiPolygon, Polygon
+from spatialdata.transformations import Identity, get_transformation
 
 from spatialdata_io.readers.merscope import (
     _dask_image_load_merscope,
+    _get_channel_names,
     _get_reader,
     _rioxarray_load_merscope,
+    merscope,
 )
 
+if TYPE_CHECKING:
+    from pathlib import Path
+
+    from _pytest.monkeypatch import MonkeyPatch
+
 
 @pytest.fixture
 def broken_rioxarray(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
@@ -30,3 +45,187 @@ def test_get_reader_honours_an_explicit_backend() -> None:
 def test_get_reader_falls_back_when_rioxarray_cannot_be_imported() -> None:
     """A `rioxarray` that is installed but raises on import must not select the rioxarray backend."""
     assert _get_reader(None) is _dask_image_load_merscope
+
+
+GENES = ["GeneA", "GeneB", "GeneC"]
+BLANKS = ["Blank-1", "Blank-2"]
+CELL_IDS = ["1001", "1002", "1003", "1004"]
+
+
+def _square(cx: float, cy: float, r: float = 1.0) -> Polygon:
+    return Polygon([(cx - r, cy - r), (cx + r, cy - r), (cx + r, cy + r), (cx - r, cy + r)])
+
+
+def _write_dataset(
+    root: Path,
+    *,
+    with_images_transform: bool = True,
+    polygon_geometry: bool = False,
+) -> Path:
+    """Write a tiny standard-layout MERSCOPE dataset and return the region directory."""
+    region = root / "region0"
+    region.mkdir(parents=True)
+
+    rng = np.random.default_rng(0)
+    xs = [10.0 * i for i in range(len(CELL_IDS))]
+    ys = [5.0 * i for i in range(len(CELL_IDS))]
+
+    # counts: cells x (genes + blanks)
+    counts = pd.DataFrame(
+        rng.integers(0, 5, size=(len(CELL_IDS), len(GENES) + len(BLANKS))),
+        index=pd.Index(CELL_IDS, name="cell"),
+        columns=GENES + BLANKS,
+    )
+    counts.to_csv(region / "cell_by_gene.csv")
+
+    # metadata: same cell order as the counts (as Vizgen writes them)
+    meta = pd.DataFrame(
+        {"center_x": xs, "center_y": ys, "volume": 1.0},
+        index=pd.Index(CELL_IDS, name="EntityID"),
+    )
+    meta.to_csv(region / "cell_metadata.csv")
+
+    # transcripts: every gene appears; include the junk/local columns a real file carries
+    tx = pd.DataFrame(
+        {
+            "barcode_id": range(len(GENES) * 2),
+            "global_x": np.tile(xs[: len(GENES)], 2),
+            "global_y": np.tile(ys[: len(GENES)], 2),
+            "global_z": 0.0,
+            "x": 0.0,
+            "y": 0.0,
+            "fov": 0,
+            "gene": GENES + GENES,
+            "transcript_id": [f"t{i}" for i in range(len(GENES) * 2)],
+            "cell_id": -1,
+        }
+    )
+    tx.to_csv(region / "detected_transcripts.csv")
+
+    # boundaries parquet with a capital "Geometry" column (as Vizgen VPT writes)
+    import geopandas as gpd
+
+    geoms = []
+    for cx, cy in zip(xs, ys, strict=True):
+        sq = _square(cx, cy)
+        geoms.append(sq if polygon_geometry else MultiPolygon([sq]))
+    gdf = gpd.GeoDataFrame({"EntityID": CELL_IDS, "ZIndex": 0}, geometry=geoms).rename_geometry("Geometry")
+    gdf.to_parquet(region / "cell_boundaries.parquet")
+
+    if with_images_transform:
+        images = region / "images"
+        images.mkdir()
+        (images / "micron_to_mosaic_pixel_transform.csv").write_text("1.0 0.0 0.0\n0.0 1.0 0.0\n0.0 0.0 1.0\n")
+
+    return region
+
+
+def test_table_is_sparse_and_values_aligned(tmp_path: Path) -> None:
+    region = _write_dataset(tmp_path)
+    sdata = merscope(region, mosaic_images=False)
+    adata = sdata["table"]
+
+    # X must be sparse (like the other scverse readers), not a dense int64 array
+    assert issparse(adata.X), f"table X should be sparse, got {type(adata.X)}"
+
+    # values / var names stay correct (reader preserves the count-matrix row order)
+    counts = pd.read_csv(region / "cell_by_gene.csv", index_col=0)
+    np.testing.assert_array_equal(adata.to_df()[GENES].to_numpy(), counts[GENES].to_numpy())
+    assert list(adata.var_names) == GENES
+    assert adata.obsm["blank"].shape[1] == len(BLANKS)
+
+
+def test_reading_does_not_warn_about_deprecated_dtype(tmp_path: Path) -> None:
+    region = _write_dataset(tmp_path)
+    with warnings.catch_warnings(record=True) as caught:
+        warnings.simplefilter("always")
+        merscope(region, mosaic_images=False, transcripts=False, cells_boundaries=False)
+    assert not any("dtype argument is deprecated" in str(w.message) for w in caught)
+
+
+def test_transcript_gene_categories_known_across_partitions(tmp_path: Path, monkeypatch: MonkeyPatch) -> None:
+    # A gene confined to a later dask partition must survive. Force >=2 partitions by shrinking the
+    # reader's dd.read_csv blocksize; "GeneC" only appears in the tail, so the first partition does not
+    # contain it -> on the unfixed reader its transcripts become NaN.
+    import dask.dataframe as dd
+
+    region = _write_dataset(tmp_path)
+    genes = ["GeneA", "GeneB"] * 400 + ["GeneC"] * 20
+    tx = pd.DataFrame({"global_x": 0.0, "global_y": 0.0, "global_z": 0.0, "gene": genes})
+    tx.to_csv(region / "detected_transcripts.csv", index=False)
+    read_csv = dd.read_csv
+    monkeypatch.setattr(dd, "read_csv", lambda p, **kw: read_csv(p, blocksize="2kB", **kw))
+
+    sdata = merscope(region, mosaic_images=False, cells_boundaries=False, cells_table=False)
+    gene = next(iter(sdata.points.values()))["gene"]
+    assert gene.npartitions >= 2
+    assert gene.cat.known
+    computed = gene.compute()
+    assert not computed.isna().any()
+    assert set(computed.astype(str).unique()) == set(GENES)
+
+
+def test_read_without_images_directory(tmp_path: Path) -> None:
+    # A dataset without the images/ dir must still read with the defaults (mosaic_images=True):
+    # identity transform + no channels found, not a crash.
+    region = _write_dataset(tmp_path, with_images_transform=False)
+    sdata = merscope(region)
+    assert "table" in sdata.tables
+    assert len(sdata.images) == 0
+    gdf = next(iter(sdata.shapes.values()))
+    assert isinstance(get_transformation(gdf, "global"), Identity)
+
+
+def test_vpt_dict_mode_accepts_string_paths(tmp_path: Path) -> None:
+    # dict VPT mode with string paths must not crash on `.exists()`
+    region = _write_dataset(tmp_path)
+    vpt = {
+        "cell_by_gene": str(region / "cell_by_gene.csv"),
+        "cell_metadata": str(region / "cell_metadata.csv"),
+        "cell_boundaries": str(region / "cell_boundaries.parquet"),
+    }
+    sdata = merscope(region, vpt_outputs=vpt, mosaic_images=False)
+    assert "table" in sdata.tables
+    assert len(sdata.shapes) == 1
+
+
+def test_polygon_boundaries_supported(tmp_path: Path) -> None:
+    # plain Polygon geometries must not crash the MultiPolygon normalisation
+    region = _write_dataset(tmp_path, polygon_geometry=True)
+    sdata = merscope(region, mosaic_images=False, transcripts=False, cells_table=False)
+    gdf = next(iter(sdata.shapes.values()))
+    assert (gdf.geom_type == "MultiPolygon").all()
+    assert len(gdf) == len(CELL_IDS)
+
+
+def test_invalid_boundaries_are_repaired_not_dropped(tmp_path: Path) -> None:
+    # a self-intersecting (invalid) boundary must be repaired: dropping it leaves its table row
+    # annotating a shape that does not exist
+    import geopandas as gpd
+    from shapely.geometry import Polygon
+
+    region = _write_dataset(tmp_path)
+    gdf = gpd.read_parquet(region / "cell_boundaries.parquet")
+    bowtie = Polygon([(0, 0), (4, 4), (4, 0), (0, 4)])
+    assert not bowtie.is_valid
+    gdf.loc[0, "Geometry"] = MultiPolygon([bowtie])
+    gdf.to_parquet(region / "cell_boundaries.parquet")
+
+    sdata = merscope(region, mosaic_images=False, transcripts=False)
+    shapes = next(iter(sdata.shapes.values()))
+    assert len(shapes) == len(CELL_IDS)
+    assert shapes.geometry.is_valid.all()
+    assert (shapes.geom_type == "MultiPolygon").all()
+    assert set(sdata["table"].obs["EntityID"].astype(str)) <= set(shapes.index)
+
+
+def test_channel_names_sorted_deterministic(tmp_path: Path) -> None:
+    # channel order must be deterministic (sorted), independent of set iteration order
+    images = tmp_path / "images"
+    images.mkdir()
+    stains = ["PolyT", "DAPI", "Cellbound1", "Cellbound2"]
+    for stain in stains:
+        for z in range(3):
+            (images / f"mosaic_{stain}_z{z}.tif").touch()
+    result = _get_channel_names(images)
+    assert result == sorted(stains)