Fix query of 2D/3D data with 2D/3D bounding box

.gitignore

@@ -45,3 +45,6 @@ spatialdata-sandbox
 
 # version file
 _version.py
+
+# other
+node_modules/

src/spatialdata/_core/query/spatial_query.py

@@ -61,7 +61,9 @@ def _get_bounding_box_corners_in_intrinsic_coordinates(
     target_coordinate_system
         The coordinate system the bounding box is defined in.
 
-    Returns ------- All the corners of the bounding box in the intrinsic coordinate system of the element. The shape
+    Returns
+    -------
+    All the corners of the bounding box in the intrinsic coordinate system of the element. The shape
     is (2, 4) when axes has 2 spatial dimensions, and (2, 8) when axes has 3 spatial dimensions.
 
     The axes of the intrinsic coordinate system.
@@ -73,6 +75,12 @@ def _get_bounding_box_corners_in_intrinsic_coordinates(
     # 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)
+    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(
+        axes, min_coordinate, max_coordinate, output_axes_without_c
+    )
 
     # get the coordinates of the bounding box corners
     bounding_box_corners = get_bounding_box_corners(
@@ -155,7 +163,7 @@ def _bounding_box_mask_points(
     min_coordinate: list[Number] | ArrayLike,
     max_coordinate: list[Number] | ArrayLike,
 ) -> da.Array:
-    """Compute a mask that is true for the points inside of an axis-aligned bounding box..
+    """Compute a mask that is true for the points inside an axis-aligned bounding box.
 
     Parameters
     ----------
@@ -164,23 +172,26 @@ def _bounding_box_mask_points(
     axes
         The axes that min_coordinate and max_coordinate refer to.
     min_coordinate
-        The upper left hand corner of the bounding box (i.e., minimum coordinates
-        along all dimensions).
+        The upper left hand corner of the bounding box (i.e., minimum coordinates along all dimensions).
     max_coordinate
-        The lower right hand corner of the bounding box (i.e., the maximum coordinates
-        along all dimensions
+        The lower right hand corner of the bounding box (i.e., the maximum coordinates along all dimensions).
 
     Returns
     -------
-    The mask for the points inside of the bounding box.
+    The mask for the points inside the bounding box.
     """
+    element_axes = get_axes_names(points)
     min_coordinate = _parse_list_into_array(min_coordinate)
     max_coordinate = _parse_list_into_array(max_coordinate)
     in_bounding_box_masks = []
     for axis_index, axis_name in enumerate(axes):
+        if axis_name not in element_axes:
+            continue
         min_value = min_coordinate[axis_index]
         in_bounding_box_masks.append(points[axis_name].gt(min_value).to_dask_array(lengths=True))
     for axis_index, axis_name in enumerate(axes):
+        if axis_name not in element_axes:
+            continue
         max_value = max_coordinate[axis_index]
         in_bounding_box_masks.append(points[axis_name].lt(max_value).to_dask_array(lengths=True))
     in_bounding_box_masks = da.stack(in_bounding_box_masks, axis=-1)
@@ -269,6 +280,69 @@ def _(
     return SpatialData(**new_elements, table=table)
 
 
+def _get_axes_of_tranformation(
+    element: SpatialElement, target_coordinate_system: str
+) -> tuple[ArrayLike, tuple[str, ...], tuple[str, ...]]:
+    """
+    Get the transformation (ignoring c) from the element's intrinsic coordinate system to the query coordinate space.
+
+    Note that the axes which specify the query shape are not necessarily the same as the axes that are output of the
+    transformation
+
+    Parameters
+    ----------
+    element
+    target_coordinate_system
+
+    Returns
+    -------
+    m_without_c
+        The transformation from the element's intrinsic coordinate system to the query coordinate space, without the
+        "c" axis.
+    input_axes_without_c
+        The axes of the element's intrinsic coordinate system, without the "c" axis.
+    output_axes_without_c
+        The axes of the query coordinate system, without the "c" axis.
+
+    """
+    from spatialdata.transformations import get_transformation
+
+    transform_to_query_space = get_transformation(element, to_coordinate_system=target_coordinate_system)
+    assert isinstance(transform_to_query_space, BaseTransformation)
+    m = _get_affine_for_element(element, transform_to_query_space)
+    input_axes_without_c = tuple([ax for ax in m.input_axes if ax != "c"])
+    output_axes_without_c = tuple([ax for ax in m.output_axes if ax != "c"])
+    m_without_c = m.to_affine_matrix(input_axes=input_axes_without_c, output_axes=output_axes_without_c)
+    return m_without_c, input_axes_without_c, output_axes_without_c
+
+
+def _adjust_bounding_box_to_real_axes(
+    axes: tuple[str, ...],
+    min_coordinate: ArrayLike,
+    max_coordinate: ArrayLike,
+    output_axes_without_c: tuple[str, ...],
+) -> tuple[tuple[str, ...], ArrayLike, ArrayLike]:
+    """Adjust the bounding box to the real axes of the transformation."""
+    if set(axes) != set(output_axes_without_c):
+        axes_only_in_bb = set(axes) - set(output_axes_without_c)
+        axes_only_in_output = set(output_axes_without_c) - set(axes)
+
+        # let's remove from the bounding box whose axes that are not in the output axes
+        indices_to_remove_from_bb = [axes.index(ax) for ax in axes_only_in_bb]
+        axes = tuple([ax for ax in axes if ax not in axes_only_in_bb])
+        min_coordinate = np.delete(min_coordinate, indices_to_remove_from_bb)
+        max_coordinate = np.delete(max_coordinate, indices_to_remove_from_bb)
+
+        # if there are axes in the output axes that are not in the bounding box, we need to add them to the bounding box
+        # with a huge range
+        for ax in axes_only_in_output:
+            axes = axes + (ax,)
+            M = np.finfo(np.float32).max - 1
+            min_coordinate = np.append(min_coordinate, -M)
+            max_coordinate = np.append(max_coordinate, M)
+    return axes, min_coordinate, max_coordinate
+
+
 @bounding_box_query.register(SpatialImage)
 @bounding_box_query.register(MultiscaleSpatialImage)
 def _(
@@ -299,15 +373,10 @@ def _(
         max_coordinate=max_coordinate,
     )
 
-    # get the transformation from the element's intrinsic coordinate system to the query coordinate space
-    transform_to_query_space = get_transformation(image, to_coordinate_system=target_coordinate_system)
-    assert isinstance(transform_to_query_space, BaseTransformation)
-    m = _get_affine_for_element(image, transform_to_query_space)
-    input_axes_without_c = tuple([ax for ax in m.input_axes if ax != "c"])
-    output_axes_without_c = tuple([ax for ax in m.output_axes if ax != "c"])
-    m_without_c = m.to_affine_matrix(input_axes=input_axes_without_c, output_axes=output_axes_without_c)
+    m_without_c, input_axes_without_c, output_axes_without_c = _get_axes_of_tranformation(
+        image, target_coordinate_system
+    )
     m_without_c_linear = m_without_c[:-1, :-1]
-
     transform_dimension = np.linalg.matrix_rank(m_without_c_linear)
     transform_coordinate_length = len(output_axes_without_c)
     data_dim = len(input_axes_without_c)
@@ -335,24 +404,13 @@ def _(
         error_message = (
             f"This case is not supported (data with dimension"
             f"{data_dim} but transformation with rank {transform_dimension}."
-            f"Please open a GitHub issue if you want to discuss a case."
+            f"Please open a GitHub issue if you want to discuss a use case."
         )
         raise ValueError(error_message)
 
-    if set(axes) != set(output_axes_without_c):
-        if set(axes).issubset(output_axes_without_c):
-            logger.warning(
-                f"The element has axes {output_axes_without_c}, but the query has axes {axes}. Excluding the element "
-                f"from the query result. In the future we can add support for this case. If you are interested, "
-                f"please open a GitHub issue."
-            )
-            return None
-        error_messeage = (
-            f"Invalid case. The bounding box axes are {axes},"
-            f"the spatial axes in {target_coordinate_system} are"
-            f"{output_axes_without_c}"
-        )
-        raise ValueError(error_messeage)
+    axes, min_coordinate, max_coordinate = _adjust_bounding_box_to_real_axes(
+        axes, min_coordinate, max_coordinate, output_axes_without_c
+    )
 
     spatial_transform = Affine(m_without_c, input_axes=input_axes_without_c, output_axes=output_axes_without_c)
     spatial_transform_bb_axes = Affine(
@@ -369,7 +427,7 @@ def _(
         )
     else:
         assert case == 2
-        # TODO: we need to intersect the plane in the extrinsic coordiante system with the 3D bounding box. The
+        # TODO: we need to intersect the plane in the extrinsic coordinate system with the 3D bounding box. The
         #  vertices of this polygons needs to be transformed to the intrinsic coordinate system
         raise NotImplementedError(
             "Case 2 (the transformation is embedding 2D data in the 3D space, is not "