Centralizing and simplifying transformation system in SpatialData

.gitignore

@@ -54,3 +54,29 @@ _version.py
 # benchmarking and profiling
 .asv/
 profile.speedscope.json
+
+# other
+node_modules/
+
+.mypy_cache
+.ruff_cache
+
+# Kilo and plans
+.kilo/
+plans/
+
+# test coverage
+.coverage
+htmlcov/
+
+# local hatch config
+hatch.toml
+
+# Kilo and plans
+.kilo/
+plans/
+
+# test coverage
+.coverage
+htmlcov/
+

docs/api/models_utils.md

@@ -7,7 +7,6 @@
 .. autofunction:: validate_element
 .. autodata:: SpatialElement
 .. autofunction:: get_axes_names
-.. autofunction:: get_spatial_axes
 .. autofunction:: points_geopandas_to_dask_dataframe
 .. autofunction:: points_dask_dataframe_to_geopandas
 .. autofunction:: get_channel_names

docs/api/transformations_utils.md

@@ -10,4 +10,5 @@
 .. autofunction:: get_transformation_between_landmarks
 .. autofunction:: align_elements_using_landmarks
 .. autofunction:: remove_transformations_to_coordinate_system
+.. autofunction:: get_spatial_axes
 ```

pyproject.toml

@@ -39,7 +39,8 @@ dependencies = [
   "networkx",
   "numba>=0.55",
   "numpy",
-  "ome-zarr>=0.16",
+  "ome-zarr>=0.19.2",
+  "ome-zarr-models>=1.8",
   "pandas",
   "pooch",
   "pyarrow",
@@ -51,7 +52,7 @@ dependencies = [
   "spatial-image>=1.2.3",
   "typing-extensions>=4.8",
   "universal-pathlib>=0.2.6",
-  "xarray>=2024.10",
+  "xarray>=2025.4",
   "xarray-spatial>=0.3.5",
   "zarr>=3",
 ]

src/spatialdata/__init__.py

@@ -4,7 +4,10 @@
 from importlib.metadata import version
 from typing import TYPE_CHECKING, Any
 
-import spatialdata.models._accessor
+from spatialdata.models._accessor import register_attrs_accessor_dask_df_series
+
+# Import registers Dask .attrs accessor, enabling metadata handling on Dask collections
+register_attrs_accessor_dask_df_series()
 
 __version__ = version("spatialdata")
 
@@ -59,7 +62,6 @@
     "polygon_query": "spatialdata._core.query.spatial_query",
     # _core.spatialdata
     "SpatialData": "spatialdata._core.spatialdata",
-    # _io._utils
     "get_dask_backing_files": "spatialdata._io._utils",
     # _io.format
     "SpatialDataFormatType": "spatialdata._io.format",

src/spatialdata/_core/_deepcopy.py

@@ -11,14 +11,14 @@
 from xarray import DataArray, DataTree
 
 from spatialdata._core.spatialdata import SpatialData
-from spatialdata.models._utils import SpatialElement
-from spatialdata.models.models import (
+from spatialdata.models import (
     Image2DModel,
     Image3DModel,
     Labels2DModel,
     Labels3DModel,
     PointsModel,
     RasterSchema,
+    SpatialElement,
     get_model,
 )
 

src/spatialdata/_core/centroids.py

@@ -11,9 +11,7 @@
 from xarray import DataArray, DataTree
 
 from spatialdata._core.operations.transform import transform
-from spatialdata.models import get_axes_names
-from spatialdata.models._utils import SpatialElement
-from spatialdata.models.models import Labels2DModel, Labels3DModel, PointsModel, get_model
+from spatialdata.models import Labels2DModel, Labels3DModel, PointsModel, SpatialElement, get_axes_names, get_model
 from spatialdata.transformations.operations import get_transformation
 from spatialdata.transformations.transformations import BaseTransformation
 

src/spatialdata/_core/data_extent.py

@@ -14,9 +14,7 @@
 
 from spatialdata._core.operations.transform import transform
 from spatialdata._core.spatialdata import SpatialData
-from spatialdata.models import get_axes_names
-from spatialdata.models._utils import SpatialElement
-from spatialdata.models.models import PointsModel
+from spatialdata.models import PointsModel, SpatialElement, get_axes_names
 from spatialdata.transformations.operations import get_transformation
 
 BoundingBoxDescription = dict[str, tuple[float, float]]

src/spatialdata/_core/operations/_utils.py

@@ -5,7 +5,8 @@
 from anndata import AnnData
 from xarray import DataArray, DataTree
 
-from spatialdata.models import SpatialElement, get_axes_names, get_spatial_axes
+from spatialdata.models import SpatialElement, get_axes_names
+from spatialdata.transformations import get_spatial_axes
 
 if TYPE_CHECKING:
     from spatialdata._core.spatialdata import SpatialData

src/spatialdata/_core/operations/map.py

@@ -13,7 +13,7 @@
 from xarray import DataArray, DataTree
 
 from spatialdata._types import IntArrayLike
-from spatialdata.models._utils import get_axes_names, get_channel_names, get_raster_model_from_data_dims
+from spatialdata.models import get_axes_names, get_channel_names, get_raster_model_from_data_dims
 from spatialdata.transformations import get_transformation
 
 __all__ = ["map_raster", "relabel_sequential"]

src/spatialdata/_core/operations/rasterize.py

@@ -31,7 +31,7 @@
     get_axes_names,
     get_model,
 )
-from spatialdata.models._utils import get_spatial_axes
+from spatialdata.transformations import get_spatial_axes
 from spatialdata.transformations._utils import _get_scale, compute_coordinates
 from spatialdata.transformations.operations import get_transformation, remove_transformation, set_transformation
 from spatialdata.transformations.transformations import (

src/spatialdata/_core/operations/transform.py

@@ -20,9 +20,14 @@
 from spatialdata._core.spatialdata import SpatialData
 from spatialdata._types import ArrayLike
 from spatialdata._utils import disable_dask_tune_optimization
-from spatialdata.models import SpatialElement, get_axes_names, get_model
-from spatialdata.models._utils import DEFAULT_COORDINATE_SYSTEM, get_channel_names
-from spatialdata.models.models import RasterSchema
+from spatialdata.literals import DEFAULT_COORDINATE_SYSTEM, TRANSFORM_KEY
+from spatialdata.models import (
+    RasterSchema,
+    SpatialElement,
+    get_axes_names,
+    get_channel_names,
+    get_model,
+)
 from spatialdata.transformations._utils import _get_scale, compute_coordinates, scale_radii
 
 if TYPE_CHECKING:
@@ -356,7 +361,6 @@ def _(
     schema = get_model(data)
     assert issubclass(schema, RasterSchema)
     from spatialdata.models import Image2DModel, Image3DModel, Labels2DModel, Labels3DModel
-    from spatialdata.models._utils import TRANSFORM_KEY
     from spatialdata.transformations import get_transformation, set_transformation
     from spatialdata.transformations.transformations import Identity, Sequence
 
@@ -441,7 +445,6 @@ def _(
     to_coordinate_system: str | None = None,
 ) -> DaskDataFrame:
     from spatialdata.models import PointsModel
-    from spatialdata.models._utils import TRANSFORM_KEY
     from spatialdata.transformations import Identity, get_transformation
 
     transformation, to_coordinate_system = _validate_target_coordinate_systems(
@@ -519,7 +522,6 @@ def _(
     to_coordinate_system: str | None = None,
 ) -> GeoDataFrame:
     from spatialdata.models import ShapesModel
-    from spatialdata.models._utils import TRANSFORM_KEY
     from spatialdata.transformations import Identity, get_transformation
 
     axes = get_axes_names(data)

src/spatialdata/_core/operations/vectorize.py

@@ -30,8 +30,8 @@
     SpatialElement,
     get_axes_names,
     get_model,
+    points_dask_dataframe_to_geopandas,
 )
-from spatialdata.models._utils import points_dask_dataframe_to_geopandas
 from spatialdata.transformations.operations import get_transformation
 from spatialdata.transformations.transformations import Identity
 

src/spatialdata/_core/query/spatial_query.py

@@ -21,15 +21,16 @@
 from spatialdata._types import ArrayLike, ListOrNDArrayFloating, Raster_T
 from spatialdata._utils import _parse_list_into_array
 from spatialdata.models import (
+    ATTRS_KEY,
     PointsModel,
     ShapesModel,
     SpatialElement,
     get_axes_names,
     points_dask_dataframe_to_geopandas,
     points_geopandas_to_dask_dataframe,
 )
-from spatialdata.models._utils import ValidAxis_t, get_spatial_axes
-from spatialdata.models.models import ATTRS_KEY
+from spatialdata.transformations import get_spatial_axes
+from spatialdata.transformations._graph.types import ValidAxis_t
 from spatialdata.transformations.operations import set_transformation
 from spatialdata.transformations.transformations import (
     Affine,