From 7936f48b64f9e2090d2b4e5fa180870341c1ec6a Mon Sep 17 00:00:00 2001
From: Luca Marconato <m.lucalmer@gmail.com>
Date: Tue, 28 Nov 2023 15:47:23 +0000
Subject: [PATCH 01/10] wip

---
 src/spatialdata/_core/query/spatial_query.py |  23 +--
 tests/core/query/test_spatial_query.py       | 192 ++++++++-----------
 2 files changed, 87 insertions(+), 128 deletions(-)

diff --git a/src/spatialdata/_core/query/spatial_query.py b/src/spatialdata/_core/query/spatial_query.py
index 2cbd02f3..323906c3 100644
--- a/src/spatialdata/_core/query/spatial_query.py
+++ b/src/spatialdata/_core/query/spatial_query.py
@@ -336,24 +336,21 @@ 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."
+            # TODO: hereeeeeeeeeeeeeee need to to something with input_axes_without_c and output_axes_without_c
+            pass
+        else:
+            error_message = (
+                f"Invalid case. The bounding box axes are {axes},"
+                f"the spatial axes in {target_coordinate_system} are"
+                f"{output_axes_without_c}"
             )
-            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)
+            raise ValueError(error_message)
 
     spatial_transform = Affine(m_without_c, input_axes=input_axes_without_c, output_axes=output_axes_without_c)
     spatial_transform_bb_axes = Affine(
@@ -370,7 +367,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 "
diff --git a/tests/core/query/test_spatial_query.py b/tests/core/query/test_spatial_query.py
index 6db7e904..c4b588ab 100644
--- a/tests/core/query/test_spatial_query.py
+++ b/tests/core/query/test_spatial_query.py
@@ -142,131 +142,85 @@ def test_bounding_box_points_no_points():
     assert request is None
 
 
-@pytest.mark.parametrize("n_channels", [1, 2, 3])
-def test_bounding_box_image_2d(n_channels):
-    """Apply a bounding box to a 2D image"""
-    image = np.zeros((n_channels, 10, 10))
-    # y: [5, 9], x: [0, 4] has value 1
-    image[:, 5::, 0:5] = 1
-    image_element = Image2DModel.parse(image)
-    image_element_multiscale = Image2DModel.parse(image, scale_factors=[2, 2])
+# @pytest.mark.parametrize("n_channels", [1, 2, 3])
+# @pytest.mark.parametrize("is_labels", [False, True])
+# @pytest.mark.parametrize("is_3d", [False, True])
+# @pytest.mark.parametrize("is_bb_3d", [False, True])
+@pytest.mark.parametrize("n_channels", [1])
+@pytest.mark.parametrize("is_labels", [False, True])
+@pytest.mark.parametrize("is_3d", [True])
+@pytest.mark.parametrize("is_bb_3d", [False])
+def test_bounding_box_raster(n_channels, is_labels, is_3d, is_bb_3d):
+    """Apply a bounding box to a raster element."""
+    if is_labels and n_channels > 1:
+        # labels cannot have multiple channels, let's ignore this combination of parameters
+        return
+
+    shape = (10, 10)
+    if is_3d:
+        shape = (10, ) + shape
+    if not is_labels:
+        shape = (n_channels, ) + shape
+    else:
+        shape = (1, ) + shape
+
+    image = np.zeros(shape)
+    if is_3d:
+        image[:, 5::, 0:5, 2:7] = 1
+    else:
+        image[:, 5::, 0:5] = 1
+
+    if is_labels:
+        image = np.squeeze(image, axis=0)
+
+    model = (
+        Labels3DModel
+        if is_labels and is_3d
+        else Labels2DModel
+        if is_labels
+        else Image3DModel
+        if is_3d
+        else Image2DModel
+    )
 
-    for image in [image_element, image_element_multiscale]:
-        # bounding box: y: [5, 10[, x: [0, 5[
-        image_result = bounding_box_query(
-            image,
-            axes=("y", "x"),
-            min_coordinate=np.array([5, 0]),
-            max_coordinate=np.array([10, 5]),
-            target_coordinate_system="global",
-        )
-        expected_image = np.ones((n_channels, 5, 5))  # c dimension is preserved
-        if isinstance(image, SpatialImage):
-            assert isinstance(image, SpatialImage)
-            np.testing.assert_allclose(image_result, expected_image)
-        elif isinstance(image, MultiscaleSpatialImage):
-            assert isinstance(image_result, MultiscaleSpatialImage)
-            v = image_result["scale0"].values()
-            assert len(v) == 1
-            xdata = v.__iter__().__next__()
-            np.testing.assert_allclose(xdata, expected_image)
-        else:
-            raise ValueError("Unexpected type")
+    image_element = model.parse(image)
+    image_element_multiscale = model.parse(image, scale_factors=[2, 2])
 
+    images = [image_element, image_element_multiscale]
 
-@pytest.mark.parametrize("n_channels", [1, 2, 3])
-def test_bounding_box_image_3d(n_channels):
-    """Apply a bounding box to a 3D image"""
-    image = np.zeros((n_channels, 10, 10, 10))
-    # z: [5, 9], y: [0, 4], x: [2, 6] has value 1
-    image[:, 5::, 0:5, 2:7] = 1
-    image_element = Image3DModel.parse(image)
-    image_element_multiscale = Image3DModel.parse(image, scale_factors=[2, 2])
+    for image in images:
+        if is_bb_3d:
+            _min_coordinate = np.array([5, 0, 2])
+            _max_coordinate = np.array([10, 5, 7])
+            _axes = ("z", "y", "x")
+        else:
+            _min_coordinate = np.array([5, 0])
+            _max_coordinate = np.array([10, 5])
+            _axes = ("y", "x")
 
-    for image in [image_element, image_element_multiscale]:
-        # bounding box: z: [5, 10[, y: [0, 5[, x: [2, 7[
         image_result = bounding_box_query(
             image,
-            axes=("z", "y", "x"),
-            min_coordinate=np.array([5, 0, 2]),
-            max_coordinate=np.array([10, 5, 7]),
+            axes=_axes,
+            min_coordinate=_min_coordinate,
+            max_coordinate=_max_coordinate,
             target_coordinate_system="global",
         )
-        expected_image = np.ones((n_channels, 5, 5, 5))  # c dimension is preserved
-        if isinstance(image, SpatialImage):
-            assert isinstance(image, SpatialImage)
-            np.testing.assert_allclose(image_result, expected_image)
-        elif isinstance(image, MultiscaleSpatialImage):
-            assert isinstance(image_result, MultiscaleSpatialImage)
-            v = image_result["scale0"].values()
-            assert len(v) == 1
-            xdata = v.__iter__().__next__()
-            np.testing.assert_allclose(xdata, expected_image)
-        else:
-            raise ValueError("Unexpected type")
-
-
-def test_bounding_box_labels_2d():
-    """Apply a bounding box to a 2D label image"""
-    # in this test let's try some affine transformations, we could do that also for the other tests
-    image = np.zeros((10, 10))
-    # y: [5, 9], x: [0, 4] has value 1
-    image[5::, 0:5] = 1
-    labels_element = Labels2DModel.parse(image)
-    labels_element_multiscale = Labels2DModel.parse(image, scale_factors=[2, 2])
-
-    for labels in [labels_element, labels_element_multiscale]:
-        # bounding box: y: [5, 10[, x: [0, 5[
-        labels_result = bounding_box_query(
-            labels,
-            axes=("y", "x"),
-            min_coordinate=np.array([5, 0]),
-            max_coordinate=np.array([10, 5]),
-            target_coordinate_system="global",
-        )
-        expected_image = np.ones((5, 5))
-        if isinstance(labels, SpatialImage):
-            assert isinstance(labels, SpatialImage)
-            np.testing.assert_allclose(labels_result, expected_image)
-        elif isinstance(labels, MultiscaleSpatialImage):
-            assert isinstance(labels_result, MultiscaleSpatialImage)
-            v = labels_result["scale0"].values()
-            assert len(v) == 1
-            xdata = v.__iter__().__next__()
-            np.testing.assert_allclose(xdata, expected_image)
-        else:
-            raise ValueError("Unexpected type")
 
+    expected_image = np.ones((n_channels, 5, 5, 5)) if is_3d else np.ones((n_channels, 5, 5))
+    if is_labels:
+        expected_image = np.squeeze(expected_image, axis=0)
 
-def test_bounding_box_labels_3d():
-    """Apply a bounding box to a 3D label image"""
-    image = np.zeros((10, 10, 10), dtype=int)
-    # z: [5, 9], y: [0, 4], x: [2, 6] has value 1
-    image[5::, 0:5, 2:7] = 1
-    labels_element = Labels3DModel.parse(image)
-    labels_element_multiscale = Labels3DModel.parse(image, scale_factors=[2, 2])
-
-    for labels in [labels_element, labels_element_multiscale]:
-        # bounding box: z: [5, 10[, y: [0, 5[, x: [2, 7[
-        labels_result = bounding_box_query(
-            labels,
-            axes=("z", "y", "x"),
-            min_coordinate=np.array([5, 0, 2]),
-            max_coordinate=np.array([10, 5, 7]),
-            target_coordinate_system="global",
-        )
-        expected_image = np.ones((5, 5, 5))
-        if isinstance(labels, SpatialImage):
-            assert isinstance(labels, SpatialImage)
-            np.testing.assert_allclose(labels_result, expected_image)
-        elif isinstance(labels, MultiscaleSpatialImage):
-            assert isinstance(labels_result, MultiscaleSpatialImage)
-            v = labels_result["scale0"].values()
-            assert len(v) == 1
-            xdata = v.__iter__().__next__()
-            np.testing.assert_allclose(xdata, expected_image)
-        else:
-            raise ValueError("Unexpected type")
+    if isinstance(image, SpatialImage):
+        assert isinstance(image, SpatialImage)
+        np.testing.assert_allclose(image_result, expected_image)
+    elif isinstance(image, MultiscaleSpatialImage):
+        assert isinstance(image_result, MultiscaleSpatialImage)
+        v = image_result["scale0"].values()
+        assert len(v) == 1
+        xdata = v.__iter__().__next__()
+        np.testing.assert_allclose(xdata, expected_image)
+    else:
+        raise ValueError("Unexpected type")
 
 
 # TODO: more tests can be added for spatial queries after the cases 2, 3, 4 are implemented
@@ -357,6 +311,14 @@ def test_bounding_box_filter_table():
     assert len(queried1.table) == 3
 
 
+def test_3d_bounding_box_2d_raster():
+    pass
+
+
+def test_2d_bounding_box_3d_points():
+    pass
+
+
 # ----------------- test polygon query -----------------
 def test_polygon_query_points(sdata_query_aggregation):
     sdata = sdata_query_aggregation

From a41d3d1fef99f277b1ae1e9e3e1f43bb0686e574 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Tue, 28 Nov 2023 15:55:40 +0000
Subject: [PATCH 02/10] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 tests/core/query/test_spatial_query.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/tests/core/query/test_spatial_query.py b/tests/core/query/test_spatial_query.py
index c4b588ab..a168898b 100644
--- a/tests/core/query/test_spatial_query.py
+++ b/tests/core/query/test_spatial_query.py
@@ -158,11 +158,11 @@ def test_bounding_box_raster(n_channels, is_labels, is_3d, is_bb_3d):
 
     shape = (10, 10)
     if is_3d:
-        shape = (10, ) + shape
+        shape = (10,) + shape
     if not is_labels:
-        shape = (n_channels, ) + shape
+        shape = (n_channels,) + shape
     else:
-        shape = (1, ) + shape
+        shape = (1,) + shape
 
     image = np.zeros(shape)
     if is_3d:

From af214b3eb1d0e758a25a4c33262c40b0678c2252 Mon Sep 17 00:00:00 2001
From: Luca Marconato <m.lucalmer@gmail.com>
Date: Tue, 28 Nov 2023 17:22:19 +0000
Subject: [PATCH 03/10] wip

---
 src/spatialdata/_core/query/spatial_query.py | 38 +++++++++++----
 tests/core/query/test_spatial_query.py       | 49 +++++++++++---------
 2 files changed, 56 insertions(+), 31 deletions(-)

diff --git a/src/spatialdata/_core/query/spatial_query.py b/src/spatialdata/_core/query/spatial_query.py
index 323906c3..dac8a988 100644
--- a/src/spatialdata/_core/query/spatial_query.py
+++ b/src/spatialdata/_core/query/spatial_query.py
@@ -341,16 +341,34 @@ def _(
         raise ValueError(error_message)
 
     if set(axes) != set(output_axes_without_c):
-        if set(axes).issubset(output_axes_without_c):
-            # TODO: hereeeeeeeeeeeeeee need to to something with input_axes_without_c and output_axes_without_c
-            pass
-        else:
-            error_message = (
-                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_message)
+        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)
+
+        #     if case == 1 or case == 5:
+        #         # in these cases the data is 2D or 3D and the bounding box is specified on a different number of axes (1, 2 or 3)
+        #         axes
+        #     else:
+        #         raise NotImplementedError('Cases 2, 3, 4 are not implemented yet.')
+        # else:
+        #     error_message = (
+        #         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_message)
 
     spatial_transform = Affine(m_without_c, input_axes=input_axes_without_c, output_axes=output_axes_without_c)
     spatial_transform_bb_axes = Affine(
diff --git a/tests/core/query/test_spatial_query.py b/tests/core/query/test_spatial_query.py
index c4b588ab..46fa02a5 100644
--- a/tests/core/query/test_spatial_query.py
+++ b/tests/core/query/test_spatial_query.py
@@ -2,6 +2,7 @@
 
 import numpy as np
 import pytest
+import xarray
 from anndata import AnnData
 from multiscale_spatial_image import MultiscaleSpatialImage
 from shapely import Polygon
@@ -158,21 +159,26 @@ def test_bounding_box_raster(n_channels, is_labels, is_3d, is_bb_3d):
 
     shape = (10, 10)
     if is_3d:
-        shape = (10, ) + shape
+        shape = (10,) + shape
     if not is_labels:
-        shape = (n_channels, ) + shape
+        shape = (n_channels,) + shape
     else:
-        shape = (1, ) + shape
+        shape = (1,) + shape
 
     image = np.zeros(shape)
+    axes = ['y', 'x']
     if is_3d:
-        image[:, 5::, 0:5, 2:7] = 1
+        image[:, 2:7, 5::, 0:5] = 1
+        axes = ['z'] + axes
     else:
         image[:, 5::, 0:5] = 1
 
     if is_labels:
         image = np.squeeze(image, axis=0)
+    else:
+        axes = ['c'] + axes
 
+    ximage = xarray.DataArray(image, dims=axes)
     model = (
         Labels3DModel
         if is_labels and is_3d
@@ -190,8 +196,8 @@ def test_bounding_box_raster(n_channels, is_labels, is_3d, is_bb_3d):
 
     for image in images:
         if is_bb_3d:
-            _min_coordinate = np.array([5, 0, 2])
-            _max_coordinate = np.array([10, 5, 7])
+            _min_coordinate = np.array([2, 5, 0])
+            _max_coordinate = np.array([7, 10, 5])
             _axes = ("z", "y", "x")
         else:
             _min_coordinate = np.array([5, 0])
@@ -206,21 +212,22 @@ def test_bounding_box_raster(n_channels, is_labels, is_3d, is_bb_3d):
             target_coordinate_system="global",
         )
 
-    expected_image = np.ones((n_channels, 5, 5, 5)) if is_3d else np.ones((n_channels, 5, 5))
-    if is_labels:
-        expected_image = np.squeeze(expected_image, axis=0)
-
-    if isinstance(image, SpatialImage):
-        assert isinstance(image, SpatialImage)
-        np.testing.assert_allclose(image_result, expected_image)
-    elif isinstance(image, MultiscaleSpatialImage):
-        assert isinstance(image_result, MultiscaleSpatialImage)
-        v = image_result["scale0"].values()
-        assert len(v) == 1
-        xdata = v.__iter__().__next__()
-        np.testing.assert_allclose(xdata, expected_image)
-    else:
-        raise ValueError("Unexpected type")
+        expected_image = ximage
+
+        if is_labels:
+            expected_image = np.squeeze(expected_image, axis=0)
+
+        if isinstance(image, SpatialImage):
+            assert isinstance(image, SpatialImage)
+            np.testing.assert_allclose(image_result, expected_image)
+        elif isinstance(image, MultiscaleSpatialImage):
+            assert isinstance(image_result, MultiscaleSpatialImage)
+            v = image_result["scale0"].values()
+            assert len(v) == 1
+            xdata = v.__iter__().__next__()
+            np.testing.assert_allclose(xdata, expected_image)
+        else:
+            raise ValueError("Unexpected type")
 
 
 # TODO: more tests can be added for spatial queries after the cases 2, 3, 4 are implemented

From 65516af2315e6fdad46dc2d550bf680ea4c1240f Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Tue, 28 Nov 2023 17:22:51 +0000
Subject: [PATCH 04/10] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 tests/core/query/test_spatial_query.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/tests/core/query/test_spatial_query.py b/tests/core/query/test_spatial_query.py
index 46fa02a5..78b3835a 100644
--- a/tests/core/query/test_spatial_query.py
+++ b/tests/core/query/test_spatial_query.py
@@ -166,17 +166,17 @@ def test_bounding_box_raster(n_channels, is_labels, is_3d, is_bb_3d):
         shape = (1,) + shape
 
     image = np.zeros(shape)
-    axes = ['y', 'x']
+    axes = ["y", "x"]
     if is_3d:
         image[:, 2:7, 5::, 0:5] = 1
-        axes = ['z'] + axes
+        axes = ["z"] + axes
     else:
         image[:, 5::, 0:5] = 1
 
     if is_labels:
         image = np.squeeze(image, axis=0)
     else:
-        axes = ['c'] + axes
+        axes = ["c"] + axes
 
     ximage = xarray.DataArray(image, dims=axes)
     model = (

From 4e64e16d821c781a9a7a2f09f4fc6116a9eb4afa Mon Sep 17 00:00:00 2001
From: Luca Marconato <m.lucalmer@gmail.com>
Date: Sun, 7 Jan 2024 13:44:44 +0100
Subject: [PATCH 05/10] fix 2d/3d bb for raster data

---
 .gitignore                                   |  3 +++
 src/spatialdata/_core/query/spatial_query.py | 16 ++-----------
 tests/core/query/test_spatial_query.py       | 24 ++++++++------------
 3 files changed, 14 insertions(+), 29 deletions(-)

diff --git a/.gitignore b/.gitignore
index c0a79ef9..666248e4 100644
--- a/.gitignore
+++ b/.gitignore
@@ -45,3 +45,6 @@ spatialdata-sandbox
 
 # version file
 _version.py
+
+# other
+node_modules/
diff --git a/src/spatialdata/_core/query/spatial_query.py b/src/spatialdata/_core/query/spatial_query.py
index bc71c8b5..e4a05e51 100644
--- a/src/spatialdata/_core/query/spatial_query.py
+++ b/src/spatialdata/_core/query/spatial_query.py
@@ -349,26 +349,14 @@ def _(
         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
+        # 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)
 
-        #     if case == 1 or case == 5:
-        #         # in these cases the data is 2D or 3D and the bounding box is specified on a different number of axes (1, 2 or 3)
-        #         axes
-        #     else:
-        #         raise NotImplementedError('Cases 2, 3, 4 are not implemented yet.')
-        # else:
-        #     error_message = (
-        #         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_message)
-
     spatial_transform = Affine(m_without_c, input_axes=input_axes_without_c, output_axes=output_axes_without_c)
     spatial_transform_bb_axes = Affine(
         spatial_transform.to_affine_matrix(input_axes=input_axes_without_c, output_axes=axes),
diff --git a/tests/core/query/test_spatial_query.py b/tests/core/query/test_spatial_query.py
index d3577d35..7ea1da85 100644
--- a/tests/core/query/test_spatial_query.py
+++ b/tests/core/query/test_spatial_query.py
@@ -144,13 +144,10 @@ def test_bounding_box_points_no_points():
 
 
 # @pytest.mark.parametrize("n_channels", [1, 2, 3])
-# @pytest.mark.parametrize("is_labels", [False, True])
-# @pytest.mark.parametrize("is_3d", [False, True])
-# @pytest.mark.parametrize("is_bb_3d", [False, True])
-@pytest.mark.parametrize("n_channels", [1])
-@pytest.mark.parametrize("is_labels", [False, True])
-@pytest.mark.parametrize("is_3d", [True])
-@pytest.mark.parametrize("is_bb_3d", [False])
+@pytest.mark.parametrize("n_channels", [1, 2, 3])
+@pytest.mark.parametrize("is_labels", [True, False])
+@pytest.mark.parametrize("is_3d", [True, False])
+@pytest.mark.parametrize("is_bb_3d", [True, False])
 def test_bounding_box_raster(n_channels, is_labels, is_3d, is_bb_3d):
     """Apply a bounding box to a raster element."""
     if is_labels and n_channels > 1:
@@ -160,10 +157,7 @@ def test_bounding_box_raster(n_channels, is_labels, is_3d, is_bb_3d):
     shape = (10, 10)
     if is_3d:
         shape = (10,) + shape
-    if not is_labels:
-        shape = (n_channels,) + shape
-    else:
-        shape = (1,) + shape
+    shape = (n_channels,) + shape if not is_labels else (1,) + shape
 
     image = np.zeros(shape)
     axes = ["y", "x"]
@@ -212,10 +206,10 @@ def test_bounding_box_raster(n_channels, is_labels, is_3d, is_bb_3d):
             target_coordinate_system="global",
         )
 
-        expected_image = ximage
-
-        if is_labels:
-            expected_image = np.squeeze(expected_image, axis=0)
+        slices = {"y": slice(5, 10), "x": slice(0, 5)}
+        if is_bb_3d and is_3d:
+            slices["z"] = slice(2, 7)
+        expected_image = ximage.sel(**slices)
 
         if isinstance(image, SpatialImage):
             assert isinstance(image, SpatialImage)

From 255bd4e9a0fd751ad1aba52b6ef3e0d7cbaa9e8e Mon Sep 17 00:00:00 2001
From: Luca Marconato <m.lucalmer@gmail.com>
Date: Sun, 7 Jan 2024 14:47:43 +0100
Subject: [PATCH 06/10] support for 2d/3d bb for 2d/3d points

---
 src/spatialdata/_core/query/spatial_query.py | 119 ++++++++++++++-----
 tests/core/query/test_spatial_query.py       |  49 +++++---
 2 files changed, 120 insertions(+), 48 deletions(-)

diff --git a/src/spatialdata/_core/query/spatial_query.py b/src/spatialdata/_core/query/spatial_query.py
index e4a05e51..784c97d0 100644
--- a/src/spatialdata/_core/query/spatial_query.py
+++ b/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)
@@ -339,23 +408,9 @@ def _(
         )
         raise ValueError(error_message)
 
-    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)
+    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(
diff --git a/tests/core/query/test_spatial_query.py b/tests/core/query/test_spatial_query.py
index 7ea1da85..f50cbd4b 100644
--- a/tests/core/query/test_spatial_query.py
+++ b/tests/core/query/test_spatial_query.py
@@ -102,23 +102,46 @@ def test_bounding_box_request_wrong_coordinate_order():
         )
 
 
-def test_bounding_box_points():
+@pytest.mark.parametrize("is_3d", [True, False])
+@pytest.mark.parametrize("is_bb_3d", [True, False])
+def test_bounding_box_points(is_3d: bool, is_bb_3d: bool):
     """test the points bounding box_query"""
-    points_element = _make_points(np.array([[10, 10], [20, 20], [20, 30]]))
-    original_x = np.array(points_element["x"])
-    original_y = np.array(points_element["y"])
+    data_x = np.array([10, 20, 20])
+    data_y = np.array([10, 20, 30])
+    data_z = np.array([100, 200, 300])
+
+    data = np.stack((data_x, data_y), axis=1)
+    if is_3d:
+        data = np.hstack((data, data_z.reshape(-1, 1)))
+    points_element = _make_points(data)
+
+    original_x = points_element["x"]
+    original_y = points_element["y"]
+    if is_3d:
+        original_z = points_element["z"]
+
+    if is_bb_3d:
+        _min_coordinate = np.array([18, 25, 250])
+        _max_coordinate = np.array([22, 35, 350])
+        _axes = ("x", "y", "z")
+    else:
+        _min_coordinate = np.array([18, 25])
+        _max_coordinate = np.array([22, 35])
+        _axes = ("x", "y")
 
     points_result = bounding_box_query(
         points_element,
-        axes=("x", "y"),
-        min_coordinate=np.array([18, 25]),
-        max_coordinate=np.array([22, 35]),
+        axes=_axes,
+        min_coordinate=_min_coordinate,
+        max_coordinate=_max_coordinate,
         target_coordinate_system="global",
     )
 
     # Check that the correct point was selected
     np.testing.assert_allclose(points_result["x"].compute(), [20])
     np.testing.assert_allclose(points_result["y"].compute(), [30])
+    if is_3d:
+        np.testing.assert_allclose(points_result["z"].compute(), [300])
 
     # result should be valid points element
     PointsModel.validate(points_result)
@@ -126,6 +149,8 @@ def test_bounding_box_points():
     # original element should be unchanged
     np.testing.assert_allclose(points_element["x"].compute(), original_x)
     np.testing.assert_allclose(points_element["y"].compute(), original_y)
+    if is_3d:
+        np.testing.assert_allclose(points_element["z"].compute(), original_z)
 
 
 def test_bounding_box_points_no_points():
@@ -148,7 +173,7 @@ def test_bounding_box_points_no_points():
 @pytest.mark.parametrize("is_labels", [True, False])
 @pytest.mark.parametrize("is_3d", [True, False])
 @pytest.mark.parametrize("is_bb_3d", [True, False])
-def test_bounding_box_raster(n_channels, is_labels, is_3d, is_bb_3d):
+def test_bounding_box_raster(n_channels: int, is_labels: bool, is_3d: bool, is_bb_3d: bool):
     """Apply a bounding box to a raster element."""
     if is_labels and n_channels > 1:
         # labels cannot have multiple channels, let's ignore this combination of parameters
@@ -312,14 +337,6 @@ def test_bounding_box_filter_table():
     assert len(queried1.table) == 3
 
 
-def test_3d_bounding_box_2d_raster():
-    pass
-
-
-def test_2d_bounding_box_3d_points():
-    pass
-
-
 # ----------------- test polygon query -----------------
 def test_polygon_query_points(sdata_query_aggregation):
     sdata = sdata_query_aggregation

From 82b6061eab65f7249f78415c419d19dc42b0b354 Mon Sep 17 00:00:00 2001
From: Luca Marconato <m.lucalmer@gmail.com>
Date: Sun, 7 Jan 2024 14:51:25 +0100
Subject: [PATCH 07/10] better tests

---
 tests/core/query/test_spatial_query.py | 20 ++++++++++++++------
 1 file changed, 14 insertions(+), 6 deletions(-)

diff --git a/tests/core/query/test_spatial_query.py b/tests/core/query/test_spatial_query.py
index f50cbd4b..cfd83bed 100644
--- a/tests/core/query/test_spatial_query.py
+++ b/tests/core/query/test_spatial_query.py
@@ -106,9 +106,9 @@ def test_bounding_box_request_wrong_coordinate_order():
 @pytest.mark.parametrize("is_bb_3d", [True, False])
 def test_bounding_box_points(is_3d: bool, is_bb_3d: bool):
     """test the points bounding box_query"""
-    data_x = np.array([10, 20, 20])
-    data_y = np.array([10, 20, 30])
-    data_z = np.array([100, 200, 300])
+    data_x = np.array([10, 20, 20, 20])
+    data_y = np.array([10, 20, 30, 30])
+    data_z = np.array([100, 200, 200, 300])
 
     data = np.stack((data_x, data_y), axis=1)
     if is_3d:
@@ -138,10 +138,18 @@ def test_bounding_box_points(is_3d: bool, is_bb_3d: bool):
     )
 
     # Check that the correct point was selected
-    np.testing.assert_allclose(points_result["x"].compute(), [20])
-    np.testing.assert_allclose(points_result["y"].compute(), [30])
     if is_3d:
-        np.testing.assert_allclose(points_result["z"].compute(), [300])
+        if is_bb_3d:
+            np.testing.assert_allclose(points_result["x"].compute(), [20])
+            np.testing.assert_allclose(points_result["y"].compute(), [30])
+            np.testing.assert_allclose(points_result["z"].compute(), [300])
+        else:
+            np.testing.assert_allclose(points_result["x"].compute(), [20, 20])
+            np.testing.assert_allclose(points_result["y"].compute(), [30, 30])
+            np.testing.assert_allclose(points_result["z"].compute(), [200, 300])
+    else:
+        np.testing.assert_allclose(points_result["x"].compute(), [20, 20])
+        np.testing.assert_allclose(points_result["y"].compute(), [30, 30])
 
     # result should be valid points element
     PointsModel.validate(points_result)

From eecc0b9aaccc4f01558f8ec29fb4ecceafe1dd62 Mon Sep 17 00:00:00 2001
From: LucaMarconato <2664412+LucaMarconato@users.noreply.github.com>
Date: Tue, 16 Jan 2024 12:57:16 +0100
Subject: [PATCH 08/10] applied suggestions from giovp

---
 src/spatialdata/_core/query/spatial_query.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/src/spatialdata/_core/query/spatial_query.py b/src/spatialdata/_core/query/spatial_query.py
index 784c97d0..a682f65d 100644
--- a/src/spatialdata/_core/query/spatial_query.py
+++ b/src/spatialdata/_core/query/spatial_query.py
@@ -284,7 +284,7 @@ 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.
+    Get the transformation, and transformation's axes (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
@@ -292,8 +292,9 @@ def _get_axes_of_tranformation(
     Parameters
     ----------
     element
+    	SpatialData element to be transformed.
     target_coordinate_system
-
+    	The target coordinate system for the transformation.
     Returns
     -------
     m_without_c
@@ -356,7 +357,6 @@ def _(
 
     Notes
     -----
-    _____
     See https://github.com/scverse/spatialdata/pull/151 for a detailed overview of the logic of this code,
     and for the cases the comments refer to.
     """

From bedbc800804e2437121e5408be31b8a8c946d377 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Tue, 16 Jan 2024 11:57:32 +0000
Subject: [PATCH 09/10] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 src/spatialdata/_core/query/spatial_query.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/src/spatialdata/_core/query/spatial_query.py b/src/spatialdata/_core/query/spatial_query.py
index a682f65d..242c9729 100644
--- a/src/spatialdata/_core/query/spatial_query.py
+++ b/src/spatialdata/_core/query/spatial_query.py
@@ -292,9 +292,10 @@ def _get_axes_of_tranformation(
     Parameters
     ----------
     element
-    	SpatialData element to be transformed.
+        SpatialData element to be transformed.
     target_coordinate_system
-    	The target coordinate system for the transformation.
+        The target coordinate system for the transformation.
+
     Returns
     -------
     m_without_c

From d0ff8fc64a088bc3dc8fa04a26ec5bc876f6448b Mon Sep 17 00:00:00 2001
From: Luca Marconato <m.lucalmer@gmail.com>
Date: Tue, 16 Jan 2024 13:01:15 +0100
Subject: [PATCH 10/10] added comments

---
 src/spatialdata/_core/query/spatial_query.py | 19 +++++++++++++------
 1 file changed, 13 insertions(+), 6 deletions(-)

diff --git a/src/spatialdata/_core/query/spatial_query.py b/src/spatialdata/_core/query/spatial_query.py
index a682f65d..4c75c475 100644
--- a/src/spatialdata/_core/query/spatial_query.py
+++ b/src/spatialdata/_core/query/spatial_query.py
@@ -284,17 +284,19 @@ def _get_axes_of_tranformation(
     element: SpatialElement, target_coordinate_system: str
 ) -> tuple[ArrayLike, tuple[str, ...], tuple[str, ...]]:
     """
-    Get the transformation, and transformation's axes (ignoring `c`), from the element's intrinsic coordinate system to the query coordinate space.
+    Get the transformation matrix and the transformation's axes (ignoring `c`).
 
+    The transformation is the one 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
-    	SpatialData element to be transformed.
+        SpatialData element to be transformed.
     target_coordinate_system
-    	The target coordinate system for the transformation.
+        The target coordinate system for the transformation.
+
     Returns
     -------
     m_without_c
@@ -323,19 +325,24 @@ def _adjust_bounding_box_to_real_axes(
     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."""
+    """
+    Adjust the bounding box to the real axes of the transformation.
+
+    The bounding box is defined by the user and it's axes may not coincide with the 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
+        # let's remove from the bounding box whose axes that are not in the output axes (e.g. querying 2D points with a
+        # 3D bounding box)
         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
+        # with a range that includes everything (e.g. querying 3D points with a 2D bounding box)
         for ax in axes_only_in_output:
             axes = axes + (ax,)
             M = np.finfo(np.float32).max - 1