refactor bounding box

CHANGELOG.md

@@ -14,7 +14,7 @@ and this project adheres to [Semantic Versioning][].
 
 ### Minor
 
--  Relaxing `spatial-image` package requirement #616
+-   Relaxing `spatial-image` package requirement #616
 
 ## [0.2.0] - 2024-07-03
 

src/spatialdata/_core/query/spatial_query.py

@@ -2,9 +2,10 @@
 
 import warnings
 from abc import abstractmethod
+from collections.abc import Mapping
 from dataclasses import dataclass
 from functools import singledispatch
-from typing import Any, Callable
+from typing import TYPE_CHECKING, Any, Callable
 
 import dask.array as da
 import dask.dataframe as dd
@@ -49,7 +50,7 @@ def _get_bounding_box_corners_in_intrinsic_coordinates(
     min_coordinate: list[Number] | ArrayLike,
     max_coordinate: list[Number] | ArrayLike,
     target_coordinate_system: str,
-) -> tuple[ArrayLike, tuple[str, ...]]:
+) -> tuple[DataArray, tuple[str, ...]]:
     """Get all corners of a bounding box in the intrinsic coordinates of an element.
 
     Parameters
@@ -74,38 +75,60 @@ def _get_bounding_box_corners_in_intrinsic_coordinates(
 
     The axes of the intrinsic coordinate system.
     """
-    from spatialdata.transformations import get_transformation
-
     min_coordinate = _parse_list_into_array(min_coordinate)
     max_coordinate = _parse_list_into_array(max_coordinate)
-    # get the transformation from the element's intrinsic coordinate system
-    # to the query coordinate space
-    transform_to_query_space = get_transformation(element, to_coordinate_system=target_coordinate_system)
+
+    # compute the output axes of the transformation, remove c from input and output axes, return the matrix without c
+    # and then build an affine transformation from that
     m_without_c, input_axes_without_c, output_axes_without_c = _get_axes_of_tranformation(
         element, target_coordinate_system
     )
-    axes, min_coordinate, max_coordinate = _adjust_bounding_box_to_real_axes(
+    spatial_transform = Affine(m_without_c, input_axes=input_axes_without_c, output_axes=output_axes_without_c)
+
+    # we identified 5 cases (see the responsible function for details), cases 1 and 5 correspond to invertible
+    # transformations; we focus on them
+    m_without_c_linear = m_without_c[:-1, :-1]
+    _ = _get_case_of_bounding_box_query(m_without_c_linear, input_axes_without_c, output_axes_without_c)
+
+    # adjust the bounding box to the real axes, dropping or adding eventually mismatching axes; the order of the axes is
+    # not adjusted
+    axes_adjusted, min_coordinate, max_coordinate = _adjust_bounding_box_to_real_axes(
         axes, min_coordinate, max_coordinate, output_axes_without_c
     )
+    if set(axes_adjusted) != set(output_axes_without_c):
+        raise ValueError("The axes of the bounding box must match the axes of the transformation.")
+
+    # axes, input_axes_without_c = input_axes_without_c, axes
+
+    # let's get the bounding box corners and inverse transform then to the intrinsic coordinate system; since we are
+    # in case 1 or 5, the transformation is invertible
+    spatial_transform_bb_axes = Affine(
+        spatial_transform.to_affine_matrix(input_axes=input_axes_without_c, output_axes=axes_adjusted),
+        input_axes=input_axes_without_c,
+        output_axes=axes_adjusted,
+    )
 
-    # get the coordinates of the bounding box corners
     bounding_box_corners = get_bounding_box_corners(
         min_coordinate=min_coordinate,
         max_coordinate=max_coordinate,
-        axes=axes,
-    ).data
-
-    # transform the coordinates to the intrinsic coordinate system
-    intrinsic_axes = get_axes_names(element)
-    transform_to_intrinsic = transform_to_query_space.inverse().to_affine_matrix(  # type: ignore[union-attr]
-        input_axes=axes, output_axes=intrinsic_axes
+        axes=axes_adjusted,
     )
-    rotation_matrix = transform_to_intrinsic[0:-1, 0:-1]
-    translation = transform_to_intrinsic[0:-1, -1]
 
-    intrinsic_bounding_box_corners = bounding_box_corners @ rotation_matrix.T + translation
+    inverse = spatial_transform_bb_axes.inverse()
+    if not isinstance(inverse, Affine):
+        raise RuntimeError("This should not happen")
+    rotation_matrix = inverse.matrix[0:-1, 0:-1]
+    translation = inverse.matrix[0:-1, -1]
+
+    intrinsic_bounding_box_corners = bounding_box_corners.data @ rotation_matrix.T + translation
 
-    return intrinsic_bounding_box_corners, intrinsic_axes
+    return (
+        DataArray(
+            intrinsic_bounding_box_corners,
+            coords={"corner": range(len(bounding_box_corners)), "axis": list(inverse.output_axes)},
+        ),
+        input_axes_without_c,
+    )
 
 
 def _get_polygon_in_intrinsic_coordinates(
@@ -227,6 +250,11 @@ def _adjust_bounding_box_to_real_axes(
             M = np.finfo(np.float32).max - 1
             min_coordinate = np.append(min_coordinate, -M)
             max_coordinate = np.append(max_coordinate, M)
+    else:
+        indices = [axes_bb.index(ax) for ax in axes_out_without_c]
+        min_coordinate = min_coordinate[np.array(indices)]
+        max_coordinate = max_coordinate[np.array(indices)]
+        axes_bb = axes_out_without_c
     return axes_bb, min_coordinate, max_coordinate
 
 
@@ -407,6 +435,7 @@ def bounding_box_query(
     min_coordinate: list[Number] | ArrayLike,
     max_coordinate: list[Number] | ArrayLike,
     target_coordinate_system: str,
+    return_request_only: bool = False,
     filter_table: bool = True,
     **kwargs: Any,
 ) -> SpatialElement | SpatialData | None:
@@ -426,6 +455,9 @@ def bounding_box_query(
     filter_table
         If `True`, the table is filtered to only contain rows that are annotating regions
         contained within the bounding box.
+    return_request_only
+        If `True`, the function returns the bounding box coordinates in the target coordinate system.
+        Only valid with `DataArray` and `DataTree` elements.
 
     Returns
     -------
@@ -472,7 +504,8 @@ def _(
     min_coordinate: list[Number] | ArrayLike,
     max_coordinate: list[Number] | ArrayLike,
     target_coordinate_system: str,
-) -> DataArray | DataTree | None:
+    return_request_only: bool = False,
+) -> DataArray | DataTree | Mapping[str, slice] | None:
     """Implement bounding box query for Spatialdata supported DataArray.
 
     Notes
@@ -493,48 +526,11 @@ def _(
         max_coordinate=max_coordinate,
     )
 
-    # compute the output axes of the transformation, remove c from input and output axes, return the matrix without c
-    # and then build an affine transformation from that
-    m_without_c, input_axes_without_c, output_axes_without_c = _get_axes_of_tranformation(
-        image, target_coordinate_system
-    )
-    spatial_transform = Affine(m_without_c, input_axes=input_axes_without_c, output_axes=output_axes_without_c)
-
-    # we identified 5 cases (see the responsible function for details), cases 1 and 5 correspond to invertible
-    # transformations; we focus on them
-    m_without_c_linear = m_without_c[:-1, :-1]
-    case = _get_case_of_bounding_box_query(m_without_c_linear, input_axes_without_c, output_axes_without_c)
-    assert case in [1, 5]
-
-    # adjust the bounding box to the real axes, dropping or adding eventually mismatching axes; the order of the axes is
-    # not adjusted
-    axes, min_coordinate, max_coordinate = _adjust_bounding_box_to_real_axes(
-        axes, min_coordinate, max_coordinate, output_axes_without_c
-    )
-    assert set(axes) == set(output_axes_without_c)
-
-    # since the order of the axes is arbitrary, let's adjust the affine transformation without c to match those axes
-    spatial_transform_bb_axes = Affine(
-        spatial_transform.to_affine_matrix(input_axes=input_axes_without_c, output_axes=axes),
-        input_axes=input_axes_without_c,
-        output_axes=axes,
-    )
-
-    # let's get the bounding box corners and inverse transform then to the intrinsic coordinate system; since we are
-    # in case 1 or 5, the transformation is invertible
-    bounding_box_corners = get_bounding_box_corners(
-        min_coordinate=min_coordinate,
-        max_coordinate=max_coordinate,
-        axes=axes,
-    )
-    inverse = spatial_transform_bb_axes.inverse()
-    assert isinstance(inverse, Affine)
-    rotation_matrix = inverse.matrix[0:-1, 0:-1]
-    translation = inverse.matrix[0:-1, -1]
-    intrinsic_bounding_box_corners = DataArray(
-        bounding_box_corners.data @ rotation_matrix.T + translation,
-        coords={"corner": range(len(bounding_box_corners)), "axis": list(inverse.output_axes)},
+    intrinsic_bounding_box_corners, axes = _get_bounding_box_corners_in_intrinsic_coordinates(
+        image, axes, min_coordinate, max_coordinate, target_coordinate_system
     )
+    if TYPE_CHECKING:
+        assert isinstance(intrinsic_bounding_box_corners, DataArray)
 
     # build the request: now that we have the bounding box corners in the intrinsic coordinate system, we can use them
     # to build the request to query the raster data using the xarray APIs
@@ -555,6 +551,9 @@ def _(
         else:
             translation_vector.append(0)
 
+    if return_request_only:
+        return selection
+
     # query the data
     query_result = image.sel(selection)
     if isinstance(image, DataArray):
@@ -652,6 +651,7 @@ def _(
         max_coordinate=max_coordinate,
         target_coordinate_system=target_coordinate_system,
     )
+    intrinsic_bounding_box_corners = intrinsic_bounding_box_corners.data
     min_coordinate_intrinsic = intrinsic_bounding_box_corners.min(axis=0)
     max_coordinate_intrinsic = intrinsic_bounding_box_corners.max(axis=0)
 
@@ -722,14 +722,14 @@ def _(
     )
 
     # get the four corners of the bounding box
-    (intrinsic_bounding_box_corners, intrinsic_axes) = _get_bounding_box_corners_in_intrinsic_coordinates(
+    (intrinsic_bounding_box_corners, _) = _get_bounding_box_corners_in_intrinsic_coordinates(
         element=polygons,
         axes=axes,
         min_coordinate=min_coordinate,
         max_coordinate=max_coordinate,
         target_coordinate_system=target_coordinate_system,
     )
-
+    intrinsic_bounding_box_corners = intrinsic_bounding_box_corners.data
     bounding_box_non_axes_aligned = Polygon(intrinsic_bounding_box_corners)
     indices = polygons.geometry.intersects(bounding_box_non_axes_aligned)
     queried = polygons[indices]
@@ -841,6 +841,7 @@ def _(
     image: DataArray | DataTree,
     polygon: Polygon | MultiPolygon,
     target_coordinate_system: str,
+    return_request_only: bool = False,
     **kwargs: Any,
 ) -> DataArray | DataTree | None:
     _check_deprecated_kwargs(kwargs)
@@ -852,6 +853,7 @@ def _(
         max_coordinate=[max_x, max_y],
         axes=("x", "y"),
         target_coordinate_system=target_coordinate_system,
+        return_request_only=return_request_only,
     )
 
 

src/spatialdata/dataloader/datasets.py

@@ -7,6 +7,7 @@
 from types import MappingProxyType
 from typing import Any, Callable
 
+import anndata as ad
 import numpy as np
 import pandas as pd
 from anndata import AnnData
@@ -276,7 +277,7 @@ def _preprocess(
         self.dataset_index = pd.concat(index_df).reset_index(drop=True)
         assert len(self.tiles_coords) == len(self.dataset_index)
         if table_name:
-            self.dataset_table = AnnData.concatenate(*tables_l)
+            self.dataset_table = ad.concat(*tables_l)
             assert len(self.tiles_coords) == len(self.dataset_table)
 
         dims_ = set(chain(*dims_l))