Merge pull request #282 from melonora/fix_blobs

Fix blobs

Author: LucaMarconato

src/spatialdata/datasets.py

@@ -1,5 +1,5 @@
 """SpatialData datasets."""
-from typing import Optional, Union
+from typing import Any, Optional, Union
 
 import numpy as np
 import pandas as pd
@@ -22,11 +22,14 @@
     ShapesModel,
     TableModel,
 )
+from spatialdata.transformations import Identity
 
 __all__ = ["blobs", "raccoon"]
 
 
-def blobs(length: int = 512, n_points: int = 200, n_shapes: int = 5) -> SpatialData:
+def blobs(
+    length: int = 512, n_points: int = 200, n_shapes: int = 5, extra_coord_system: Optional[str] = None
+) -> SpatialData:
     """
     Blobs dataset.
 
@@ -39,13 +42,18 @@ def blobs(length: int = 512, n_points: int = 200, n_shapes: int = 5) -> SpatialD
     n_shapes
         Number of max shapes to generate.
         At most, as if overlapping they will be discarded
+    extra_coord_system
+            Extra coordinate space on top of the standard global coordinate space. Will have only identity transform.
+
 
     Returns
     -------
     SpatialData
         SpatialData object with blobs dataset.
     """
-    return BlobsDataset(length=length, n_points=n_points, n_shapes=n_shapes).blobs()
+    return BlobsDataset(
+        length=length, n_points=n_points, n_shapes=n_shapes, extra_coord_system=extra_coord_system
+    ).blobs()
 
 
 def raccoon() -> SpatialData:
@@ -75,7 +83,9 @@ def raccoon(
 class BlobsDataset:
     """Blobs dataset."""
 
-    def __init__(self, length: int = 512, n_points: int = 200, n_shapes: int = 5) -> None:
+    def __init__(
+        self, length: int = 512, n_points: int = 200, n_shapes: int = 5, extra_coord_system: Optional[str] = None
+    ) -> None:
         """
         Blobs dataset.
 
@@ -88,23 +98,28 @@ def __init__(self, length: int = 512, n_points: int = 200, n_shapes: int = 5) ->
         n_shapes
             Number of max shapes to generate.
             At most, as if overlapping they will be discarded
+        extra_coord_system
+            Extra coordinate space on top of the standard global coordinate space. Will have only identity transform.
         """
         self.length = length
         self.n_points = n_points
         self.n_shapes = n_shapes
+        self.transformations = {"global": Identity()}
+        if extra_coord_system:
+            self.transformations[extra_coord_system] = Identity()
 
     def blobs(
         self,
     ) -> SpatialData:
         """Blobs dataset."""
-        image = self._image_blobs(self.length)
-        multiscale_image = self._image_blobs(self.length, multiscale=True)
-        labels = self._labels_blobs(self.length)
-        multiscale_labels = self._labels_blobs(self.length, multiscale=True)
-        points = self._points_blobs(self.length, self.n_points)
-        circles = self._circles_blobs(self.length, self.n_shapes)
-        polygons = self._polygons_blobs(self.length, self.n_shapes)
-        multipolygons = self._polygons_blobs(self.length, self.n_shapes, multipolygons=True)
+        image = self._image_blobs(self.transformations, self.length)
+        multiscale_image = self._image_blobs(self.transformations, self.length, multiscale=True)
+        labels = self._labels_blobs(self.transformations, self.length)
+        multiscale_labels = self._labels_blobs(self.transformations, self.length, multiscale=True)
+        points = self._points_blobs(self.transformations, self.length, self.n_points)
+        circles = self._circles_blobs(self.transformations, self.length, self.n_shapes)
+        polygons = self._polygons_blobs(self.transformations, self.length, self.n_shapes)
+        multipolygons = self._polygons_blobs(self.transformations, self.length, self.n_shapes, multipolygons=True)
         adata = aggregate(image, labels)
         adata.obs["region"] = pd.Categorical(["blobs_labels"] * len(adata))
         adata.obs["instance_id"] = adata.obs_names.astype(int)
@@ -118,7 +133,12 @@ def blobs(
             table=table,
         )
 
-    def _image_blobs(self, length: int = 512, multiscale: bool = False) -> Union[SpatialImage, MultiscaleSpatialImage]:
+    def _image_blobs(
+        self,
+        transformations: Optional[dict[str, Any]] = None,
+        length: int = 512,
+        multiscale: bool = False,
+    ) -> Union[SpatialImage, MultiscaleSpatialImage]:
         masks = []
         for i in range(3):
             mask = self._generate_blobs(length=length, seed=i)
@@ -128,10 +148,12 @@ def _image_blobs(self, length: int = 512, multiscale: bool = False) -> Union[Spa
         x = np.stack(masks, axis=0)
         dims = ["c", "y", "x"]
         if not multiscale:
-            return Image2DModel.parse(x, dims=dims)
-        return Image2DModel.parse(x, dims=dims, scale_factors=[2, 2])
+            return Image2DModel.parse(x, transformations=transformations, dims=dims)
+        return Image2DModel.parse(x, transformations=transformations, dims=dims, scale_factors=[2, 2])
 
-    def _labels_blobs(self, length: int = 512, multiscale: bool = False) -> Union[SpatialImage, MultiscaleSpatialImage]:
+    def _labels_blobs(
+        self, transformations: Optional[dict[str, Any]] = None, length: int = 512, multiscale: bool = False
+    ) -> Union[SpatialImage, MultiscaleSpatialImage]:
         """Create a 2D labels."""
         from scipy.ndimage import watershed_ift
 
@@ -155,8 +177,8 @@ def _labels_blobs(self, length: int = 512, multiscale: bool = False) -> Union[Sp
                 out[out == val[idx]] = i
         dims = ["y", "x"]
         if not multiscale:
-            return Labels2DModel.parse(out, dims=dims)
-        return Labels2DModel.parse(out, dims=dims, scale_factors=[2, 2])
+            return Labels2DModel.parse(out, transformations=transformations, dims=dims)
+        return Labels2DModel.parse(out, transformations=transformations, dims=dims, scale_factors=[2, 2])
 
     def _generate_blobs(self, length: int = 512, seed: Optional[int] = None) -> ArrayLike:
         from scipy.ndimage import gaussian_filter
@@ -171,21 +193,27 @@ def _generate_blobs(self, length: int = 512, seed: Optional[int] = None) -> Arra
         mask = gaussian_filter(mask, sigma=0.25 * length * 0.1)
         return mask
 
-    def _points_blobs(self, length: int = 512, n_points: int = 200) -> DaskDataFrame:
+    def _points_blobs(
+        self, transformations: Optional[dict[str, Any]] = None, length: int = 512, n_points: int = 200
+    ) -> DaskDataFrame:
         rng = default_rng(42)
-        arr = rng.integers(10, length - 10, size=(n_points, 2)).astype(np.int_)
+        arr = rng.integers(10, length - 10, size=(n_points, 2)).astype(np.int64)
         # randomly assign some values from v to the points
-        points_assignment0 = rng.integers(0, 10, size=arr.shape[0]).astype(np.int_)
+        points_assignment0 = rng.integers(0, 10, size=arr.shape[0]).astype(np.int64)
         genes = rng.choice(["a", "b"], size=arr.shape[0])
         annotation = pd.DataFrame(
             {
                 "genes": genes,
                 "instance_id": points_assignment0,
             },
         )
-        return PointsModel.parse(arr, annotation=annotation, feature_key="genes", instance_key="instance_id")
+        return PointsModel.parse(
+            arr, transformations=transformations, annotation=annotation, feature_key="genes", instance_key="instance_id"
+        )
 
-    def _circles_blobs(self, length: int = 512, n_shapes: int = 5) -> GeoDataFrame:
+    def _circles_blobs(
+        self, transformations: Optional[dict[str, Any]] = None, length: int = 512, n_shapes: int = 5
+    ) -> GeoDataFrame:
         midpoint = length // 2
         halfmidpoint = midpoint // 2
         radius = length // 10
@@ -195,9 +223,15 @@ def _circles_blobs(self, length: int = 512, n_shapes: int = 5) -> GeoDataFrame:
                 "radius": radius,
             }
         )
-        return ShapesModel.parse(circles)
+        return ShapesModel.parse(circles, transformations=transformations)
 
-    def _polygons_blobs(self, length: int = 512, n_shapes: int = 5, multipolygons: bool = False) -> GeoDataFrame:
+    def _polygons_blobs(
+        self,
+        transformations: Optional[dict[str, Any]] = None,
+        length: int = 512,
+        n_shapes: int = 5,
+        multipolygons: bool = False,
+    ) -> GeoDataFrame:
         midpoint = length // 2
         halfmidpoint = midpoint // 2
         poly = GeoDataFrame(
@@ -207,7 +241,7 @@ def _polygons_blobs(self, length: int = 512, n_shapes: int = 5, multipolygons: b
                 )
             }
         )
-        return ShapesModel.parse(poly)
+        return ShapesModel.parse(poly, transformations=transformations)
 
     # function that generates random shapely polygons given a bounding box
     def _generate_random_polygons(

tests/datasets/test_datasets.py

@@ -2,7 +2,8 @@
 
 
 def test_datasets() -> None:
-    sdata_blobs = blobs()
+    extra_cs = "test"
+    sdata_blobs = blobs(extra_coord_system=extra_cs)
 
     assert len(sdata_blobs.table) == 26
     assert len(sdata_blobs.shapes["blobs_circles"]) == 5
@@ -13,6 +14,7 @@ def test_datasets() -> None:
     assert len(sdata_blobs.images["blobs_multiscale_image"]) == 3
     assert sdata_blobs.labels["blobs_labels"].shape == (512, 512)
     assert len(sdata_blobs.labels["blobs_multiscale_labels"]) == 3
+    assert extra_cs in sdata_blobs.coordinate_systems
     # this catches this bug: https://github.com/scverse/spatialdata/issues/269
     _ = str(sdata_blobs)