Basic handling of filter_bbox in CubeMetadata, down to SpatialDimension

refs #827, #949

Author: soxofaan

openeo/metadata/__init__.py

@@ -22,8 +22,8 @@
 
 from openeo.api.process import Parameter
 from openeo.internal.jupyter import render_component
-from openeo.metadata._extents import _EmptyExtent, _NonConcreteExtent
-from openeo.util import Rfc3339, deep_get
+from openeo.metadata._extents import NumericalInterval, _EmptyExtent, _NonConcreteExtent
+from openeo.util import BBoxDict, Rfc3339, deep_get
 from openeo.utils.datetime import DateTimeInterval
 from openeo.utils.normalize import normalize_resample_resolution, unique
 
@@ -88,15 +88,17 @@ def __init__(
         *,
         # TODO: axis should be required (without None as default), but currently optional (with fallback guess) for backwards compatibility.
         axis: Literal["x", "y"] | None = None,
-        extent: Union[Tuple[float, float], List[float]],
+        extent: Tuple[float, float] | List[float] | NumericalInterval | _NonConcreteExtent | _EmptyExtent,
         crs: Union[str, int, dict] = DEFAULT_CRS,
         step=None,
     ):
         """
         @param name:
         @param axis: Axis of the spatial raster dimension ("x" or "y")
         @param extent:
-        @param crs:
+        @param crs: The spatial reference system for the data and datacube metadata (extent, values and step),
+            specified as numerical EPSG code, WKT2 (ISO 19162) string or PROJJSON object.
+            Defaults to EPSG code 4326.
         @param step: The space between the values. Use null for irregularly spaced steps.
         """
         super().__init__(type="spatial", name=name)
@@ -110,6 +112,64 @@ def __init__(
     def rename(self, name) -> Dimension:
         return SpatialDimension(name=name, axis=self.axis, extent=self.extent, crs=self.crs, step=self.step)
 
+    def filter_extent(
+        self,
+        extent: (
+            Tuple[float | int, float | int]
+            | List[float | int]
+            | NumericalInterval
+            | _EmptyExtent
+            | _NonConcreteExtent
+            | None
+        ),
+    ):
+        if (this_extent := NumericalInterval.try_from(self.extent)) and (
+            other_extent := NumericalInterval.try_from(extent)
+        ):
+            extent = this_extent.intersection(other_extent)
+        elif isinstance(self.extent, _EmptyExtent) or isinstance(extent, _EmptyExtent):
+            extent = _EmptyExtent()
+        elif extent is None:
+            extent = self.extent
+        else:
+            # As lenient fallback, handle as non-concrete extent
+            extent = _NonConcreteExtent()
+        return SpatialDimension(name=self.name, axis=self.axis, extent=extent, crs=self.crs, step=self.step)
+
+
+class SpatialDimensionsXY(NamedTuple):
+    """
+    Simple container class for a pair of spatial dimensions (x and y).
+    """
+
+    x: SpatialDimension
+    y: SpatialDimension
+
+    def filter_bbox(self, extent: Any) -> SpatialDimensionsXY:
+        # TODO: support more types and forms of the extent (geometries, etc.) ... e.g. using the BBoxDict factories?
+        if isinstance(extent, dict) and {"west", "east", "south", "north"}.issubset(extent):
+            bbox = BBoxDict.from_dict(extent)
+            bbox_crs = bbox.get("crs", SpatialDimension.DEFAULT_CRS)
+            if bbox_crs == self.x.crs and bbox_crs == self.y.crs:
+                # No reprojection needed
+                extent_x = bbox["west"], bbox["east"]
+                extent_y = bbox["south"], bbox["north"]
+            else:
+                # TODO: reproject
+                extent_x, extent_y = _NonConcreteExtent(), _NonConcreteExtent()
+        elif isinstance(extent, _EmptyExtent):
+            extent_x, extent_y = _EmptyExtent(), _EmptyExtent()
+        elif isinstance(extent, _NonConcreteExtent):
+            extent_x, extent_y = _NonConcreteExtent(), _NonConcreteExtent()
+        else:
+            # As lenient fallback, handle as non-concrete extent
+            extent_x, extent_y = _NonConcreteExtent(), _NonConcreteExtent()
+
+        return SpatialDimensionsXY(
+            x=self.x.filter_extent(extent=extent_x),
+            y=self.y.filter_extent(extent=extent_y),
+        )
+
 
 class TemporalDimension(Dimension):
     def __init__(
@@ -330,23 +390,24 @@ def __init__(self, dimensions: Optional[List[Dimension]] = None):
         #       but most usage actually assumes it is a list that can be iterated over.
         #       Can we handle this more consistently and less error-prone?
         self._dimensions: Union[List[Dimension], None] = dimensions
-        self._band_dimension = None
-        self._temporal_dimension = None
+        self._band_dimension: BandDimension | None = None
+        self._temporal_dimension: TemporalDimension | None = None
+        self._spatial_dimensions: SpatialDimensionsXY | None = None
 
         if dimensions is not None:
-            for dim in self._dimensions:
-                # TODO: here we blindly pick last bands or temporal dimension if multiple. Let user choose?
-                # TODO: add spatial dimension handling?
-                if dim.type == "bands":
-                    if isinstance(dim, BandDimension):
-                        self._band_dimension = dim
-                    else:
-                        raise MetadataException("Invalid band dimension {d!r}".format(d=dim))
-                if dim.type == "temporal":
-                    if isinstance(dim, TemporalDimension):
-                        self._temporal_dimension = dim
-                    else:
-                        raise MetadataException("Invalid temporal dimension {d!r}".format(d=dim))
+            if band_dims := [d for d in dimensions if isinstance(d, BandDimension)]:
+                if len(band_dims) > 1:
+                    _log.warning(f"Multiple band dimensions, arbitrarily picking the first one: {band_dims=}")
+                self._band_dimension = band_dims[0]
+            if temporal_dims := [d for d in dimensions if isinstance(d, TemporalDimension)]:
+                if len(temporal_dims) > 1:
+                    _log.warning(f"Multiple temporal dimensions, arbitrarily picking the first one: {temporal_dims=}")
+                self._temporal_dimension = temporal_dims[0]
+            if spatial_dims := sorted((d for d in dimensions if isinstance(d, SpatialDimension)), key=lambda d: d.axis):
+                if [d.axis for d in spatial_dims] == ["x", "y"]:
+                    self._spatial_dimensions = SpatialDimensionsXY(x=spatial_dims[0], y=spatial_dims[1])
+                else:
+                    _log.warning(f"Failed to detect x and y spatial dimensions: {spatial_dims=}")
 
     def __eq__(self, o: Any) -> bool:
         return isinstance(o, type(self)) and self._dimensions == o._dimensions
@@ -411,9 +472,14 @@ def temporal_dimension(self) -> TemporalDimension:
             raise MetadataException("No temporal dimension")
         return self._temporal_dimension
 
+    def has_spatial_dimensions(self) -> bool:
+        return isinstance(self._spatial_dimensions, SpatialDimensionsXY)
+
     @property
-    def spatial_dimensions(self) -> List[SpatialDimension]:
-        return [d for d in self._dimensions if isinstance(d, SpatialDimension)]
+    def spatial_dimensions(self) -> SpatialDimensionsXY:
+        if not self.has_spatial_dimensions():
+            raise MetadataException("No spatial dimensions")
+        return self._spatial_dimensions
 
     def has_geometry_dimension(self):
         return any(isinstance(d, GeometryDimension) for d in self._dimensions)
@@ -453,6 +519,17 @@ def filter_temporal(self, extent: Tuple[str | None, str | None] | List[str | Non
             ]
         )
 
+    def filter_bbox(self, extent: Any) -> CubeMetadata:
+        if self._dimensions and isinstance(self._spatial_dimensions, SpatialDimensionsXY):
+            spatial_dimensions: SpatialDimensionsXY = self._spatial_dimensions.filter_bbox(extent=extent)
+            dimensions = [
+                getattr(spatial_dimensions, d.axis) if isinstance(d, SpatialDimension) else d for d in self._dimensions
+            ]
+            return self._clone_and_update(dimensions=dimensions)
+
+        _log.warning("filter_bbox on cube metadata without spatial dimension. Returning unchanged metadata.")
+        return self
+
     def filter_bands(self, band_names: List[Union[int, str]]) -> CubeMetadata:
         """
         Create new `CubeMetadata` with filtered band dimension

openeo/metadata/_extents.py

@@ -1,5 +1,8 @@
 from __future__ import annotations
 
+import dataclasses
+import typing
+
 
 class _NonConcreteExtent:
     """
@@ -41,3 +44,44 @@ def __eq__(self, other) -> bool:
         if isinstance(other, _EmptyExtent):
             return True
         return NotImplemented
+
+
+class NumericalInterval(typing.NamedTuple):
+    """
+    Simple numerical interval defined by a start and end point
+    each of type float/int;
+    or None to indicate being open-ended in that direction (no limit).
+    """
+
+    start: float | int | None
+    end: float | int | None
+
+    @classmethod
+    def try_from(cls, x: typing.Any) -> NumericalInterval | None:
+        if isinstance(x, (list, tuple)) and len(x) == 2 and all(isinstance(v, (int, float)) or v is None for v in x):
+            return NumericalInterval(*x)
+        return None
+
+    def intersection(self, other: NumericalInterval) -> NumericalInterval | _EmptyExtent:
+        """
+        Compute the intersection of this interval with another interval.
+        If the intervals do not overlap, return an _EmptyExtent.
+        """
+        if self.start is None:
+            start = other.start
+        elif other.start is None:
+            start = self.start
+        else:
+            start = max(self.start, other.start)
+
+        if self.end is None:
+            end = other.end
+        elif other.end is None:
+            end = self.end
+        else:
+            end = min(self.end, other.end)
+
+        if start is not None and end is not None and end < start:
+            return _EmptyExtent()
+
+        return NumericalInterval(start, end)

openeo/rest/datacube.py

@@ -256,6 +256,8 @@ def load_collection(
             if (connection and fetch_metadata)
             else None
         )
+        if metadata and spatial_extent:
+            metadata = metadata.filter_bbox(extent=spatial_extent)
         if metadata and temporal_extent:
             metadata = metadata.filter_temporal(extent=temporal_extent)
         if bands is not None:
@@ -455,7 +457,7 @@ def load_stac(
         """
         arguments: dict = {"url": url}
         if spatial_extent:
-            arguments["spatial_extent"] = _get_geometry_argument(
+            spatial_extent = _get_geometry_argument(
                 argument=spatial_extent,
                 valid_geojson_types=["Polygon", "MultiPolygon", "Feature", "FeatureCollection"],
                 connection=connection,
@@ -465,6 +467,7 @@ def load_stac(
                 argument_name="spatial_extent",
                 process_id="load_stac",
             )
+            arguments["spatial_extent"] = spatial_extent
 
         if temporal_extent:
             temporal_extent = cls._get_temporal_extent(extent=temporal_extent)
@@ -510,6 +513,8 @@ def load_stac(
             )
             metadata = None
 
+        if metadata and spatial_extent:
+            metadata = metadata.filter_bbox(extent=spatial_extent)
         if metadata and temporal_extent:
             metadata = metadata.filter_temporal(extent=temporal_extent)
 
@@ -713,12 +718,14 @@ def filter_bbox(
             extent = {'west': west, 'east': east, 'north': north, 'south': south}
             extent.update(dict_no_none(crs=crs, base=base, height=height))
 
+        metadata = self.metadata.filter_bbox(extent=extent) if self.metadata else None
         return self.process(
-            process_id='filter_bbox',
+            process_id="filter_bbox",
             arguments={
-                'data': THIS,
-                'extent': extent
-            }
+                "data": THIS,
+                "extent": extent,
+            },
+            metadata=metadata,
         )
 
     @openeo_process
@@ -774,12 +781,14 @@ def filter_spatial(
             "FeatureCollection",
         ]
         geometries = _get_geometry_argument(geometries, valid_geojson_types=valid_geojson_types, connection=self.connection, crs=None)
+        metadata = self.metadata.filter_bbox(extent=geometries) if self.metadata else None
         return self.process(
-            process_id='filter_spatial',
+            process_id="filter_spatial",
             arguments={
-                'data': THIS,
-                'geometries': geometries
-            }
+                "data": THIS,
+                "geometries": geometries,
+            },
+            metadata=metadata,
         )
 
     @openeo_process

openeo/util.py

@@ -535,7 +535,15 @@ class BBoxDict(dict):
     .. versionadded:: 0.10.1
     """
 
-    def __init__(self, *, west: float, south: float, east: float, north: float, crs: Optional[Union[str, int]] = None):
+    def __init__(
+        self,
+        *,
+        west: float,
+        south: float,
+        east: float,
+        north: float,
+        crs: int | str | dict | None = None,
+    ):
         super().__init__(west=west, south=south, east=east, north=north)
         if crs is not None:
             self.update(crs=normalize_crs(crs))

tests/metadata/test_extents.py

@@ -1,4 +1,6 @@
-from openeo.metadata import _EmptyExtent
+import pytest
+
+from openeo.metadata import NumericalInterval, _EmptyExtent
 
 
 class TestEmptyExtent:
@@ -9,3 +11,62 @@ def test_equality(self):
     def test_none(self):
         assert _EmptyExtent() != None
         assert None != _EmptyExtent()
+
+
+class TestNumericalInterval:
+    def test_basic(self):
+        interval = NumericalInterval(3, 5.5)
+        assert interval.start == 3
+        assert interval.end == 5.5
+
+    @pytest.mark.parametrize(["start", "end"], [(None, 5), (3, None), (None, None)])
+    def test_unbounded(self, start, end):
+        interval = NumericalInterval(start, end)
+        assert interval.start == start
+        assert interval.end == end
+
+    def test_equality(self):
+        assert NumericalInterval(1, 2) == NumericalInterval(1, 2)
+        assert NumericalInterval(1, 2) != NumericalInterval(1, 3)
+        assert NumericalInterval(1, 2) != NumericalInterval(0, 2)
+        assert NumericalInterval(None, 2) != NumericalInterval(0, 2)
+        assert NumericalInterval(None, None) == NumericalInterval(None, None)
+
+    def test_try_from(self):
+        assert NumericalInterval.try_from([1, 2]) == NumericalInterval(1, 2)
+        assert NumericalInterval.try_from((None, 5)) == NumericalInterval(None, 5)
+        assert NumericalInterval.try_from([3.5, None]) == NumericalInterval(3.5, None)
+        assert NumericalInterval.try_from([None, None]) == NumericalInterval(None, None)
+        assert NumericalInterval.try_from([1]) is None
+        assert NumericalInterval.try_from([1, 2, 3]) is None
+        assert NumericalInterval.try_from("not a list") is None
+        assert NumericalInterval.try_from([1, "not a number"]) is None
+
+    @pytest.mark.parametrize(
+        ["interval1", "interval2", "expected"],
+        [
+            # Closed interval cases
+            ((3, 5), (1, 2), _EmptyExtent()),
+            ((3, 5), (1, 3.5), NumericalInterval(3, 3.5)),
+            ((3, 5), (3.5, 4.5), NumericalInterval(3.5, 4.5)),
+            ((3, 5), (4.5, 8), NumericalInterval(4.5, 5)),
+            ((3, 5), (6, 8), _EmptyExtent()),
+            # Half-open interval cases
+            ((None, 5), (1, 2), NumericalInterval(1, 2)),
+            ((None, 5), (1, 7), NumericalInterval(1, 5)),
+            ((None, 5), (6, 7), _EmptyExtent()),
+            ((3, None), (1, 2), _EmptyExtent()),
+            ((3, None), (1, 4), NumericalInterval(3, 4)),
+            ((3, None), (5, 7), NumericalInterval(5, 7)),
+            # Unbounded interval cases
+            ((None, None), (1, 2), NumericalInterval(1, 2)),
+            ((None, None), (None, 2), NumericalInterval(None, 2)),
+            ((None, None), (3, None), NumericalInterval(3, None)),
+            ((None, None), (None, None), NumericalInterval(None, None)),
+        ],
+    )
+    def test_intersection(self, interval1, interval2, expected):
+        interval1 = NumericalInterval(*interval1)
+        interval2 = NumericalInterval(*interval2)
+        assert interval1.intersection(interval2) == expected
+        assert interval2.intersection(interval1) == expected