Add option to directly pass the bounding box as a tuple

Author: Simon-van-Diepen

sarxarray/utils.py

@@ -155,16 +155,17 @@ def _compute_coherence(numerator, denominator):
     return coherence
 
 
-def crop(data: xr.Dataset | xr.DataArray, geom: sg.Polygon) -> xr.Dataset:
+def crop(data: xr.Dataset | xr.DataArray, geom: sg.Polygon | tuple) -> xr.Dataset:
     """Crop a radar image or stack of radar images to the bounding box of a polygon.
 
     Parameters
     ----------
     data: xr.Dataset | xr.DataArray
         The dataset or data array to be cropped in azimuth and range
-    geom: sg.Polygon
+    geom: sg.Polygon | tuple
         shapely.geometry.Polygon in radar coordinates of the area that should be
-        kept, in [azimuth, range] format
+        kept, in [azimuth, range] format, OR a tuple of the bounding box of the crop in
+        (min_azimuth, min_range, max_azimuth, max_range) format
 
     Returns
     -------
@@ -174,12 +175,31 @@ def crop(data: xr.Dataset | xr.DataArray, geom: sg.Polygon) -> xr.Dataset:
     Raises
     ------
     ValueError
-        If the azimuth or range coordinate does not exist in `data`
+        - If the azimuth or range coordinate does not exist in `data`
+        - If `geom` is not `shapely.geometry.Polygon` or `tuple`
+
+    AssertionError
+        - When a tuple is passed to `geom` that falls in one of three categories:
+          - The tuple does not have exactly 4 entries
+          - The minimum azimuth coordinate is larger than or equal
+               to the maximum azimuth coordinate
+          - The minimum range coordinate is larger than or equal
+               to the maximum range coordinate
     """
     if not {"azimuth", "range"}.issubset(data.dims):
         raise ValueError("The data must have azimuth and range dimensions.")
 
-    bounding_box = geom.bounds  # returns (min_az, min_r, max_az, max_r)
+    if isinstance(geom, sg.Polygon):
+        bounding_box = geom.bounds  # returns (min_az, min_r, max_az, max_r)
+    elif isinstance(geom, tuple):
+        assert len(geom) == 4, f"geom as tuple must have four entries, got {len(geom)}"
+        assert geom[0] < geom[2], f"geom Azimuth wrong: min {geom[0]} >= max {geom[2]}"
+        assert geom[1] < geom[3], f"geom Range wrong: min {geom[1]} >= max {geom[3]}"
+        bounding_box = geom
+    else:
+        raise ValueError(
+            f"geom must be tuple or shapely.geometry.Polygon, is {type(geom)}."
+        )
     data = data.sel(
         azimuth=range(int(bounding_box[0]), int(np.ceil(bounding_box[2]))+1),
         range=range(int(bounding_box[1]), int(np.ceil(bounding_box[3]))+1)

tests/test_utils.py

@@ -38,9 +38,14 @@ def crop_geometry():
         [606, 1408],
         [606, 1409],
         [603, 1409],
-        [602, 1405]
+        [601.9, 1405]
     ])
 
+# Create a bounding box for cropping
+@pytest.fixture
+def crop_geometry_bbox():
+    return (601.9, 1405, 606, 1409)
+
 # this class tests multi_look with dataarray. For testing with dataset, see
 # test_stack.py
 class TestUtilsMultiLook:
@@ -292,17 +297,59 @@ def test_complex_coherence_bad_args(
 
 
 class TestUtilsCrop:
-    def test_crop(self, synthetic_dataarray, crop_geometry):
+    def test_crop_poly(self, synthetic_dataarray, crop_geometry):
         da = synthetic_dataarray
         geom = crop_geometry
         da_crop = crop(da, geom)
         assert da_crop.azimuth.size == 6
         assert da_crop.range.size == 5
         assert da_crop.time.size == da.time.size
 
+    def test_crop_bbox(self, synthetic_dataarray, crop_geometry_bbox):
+        da = synthetic_dataarray
+        geom = crop_geometry_bbox
+        da_crop = crop(da, geom)
+        assert da_crop.azimuth.size == 6
+        assert da_crop.range.size == 5
+        assert da_crop.time.size == da.time.size
+
     def test_crop_wrong_dimname(self, synthetic_dataarray, crop_geometry):
         da = synthetic_dataarray
         da = da.rename({"azimuth": "az"})  # rename azimuth to a wrong name
         geom = crop_geometry
         with pytest.raises(ValueError):
             _ = crop(da, geom)
+
+    def test_crop_bbox_wrong_length(self, synthetic_dataarray, crop_geometry_bbox):
+        da = synthetic_dataarray
+        geom = crop_geometry_bbox[:3]
+        with pytest.raises(AssertionError):
+            _ = crop(da, geom)
+
+    def test_crop_bbox_wrong_az(self, synthetic_dataarray, crop_geometry_bbox):
+        da = synthetic_dataarray
+        geom = (  # swap min and max azimuth
+            crop_geometry_bbox[2],
+            crop_geometry_bbox[1],
+            crop_geometry_bbox[0],
+            crop_geometry_bbox[3]
+        )
+        with pytest.raises(AssertionError):
+            _ = crop(da, geom)
+
+    def test_crop_bbox_wrong_r(self, synthetic_dataarray, crop_geometry_bbox):
+        da = synthetic_dataarray
+        geom = (  # swap min and max range
+            crop_geometry_bbox[0],
+            crop_geometry_bbox[3],
+            crop_geometry_bbox[2],
+            crop_geometry_bbox[1]
+        )
+        with pytest.raises(AssertionError):
+            _ = crop(da, geom)
+
+    def test_crop_bbox_wrong_geom_type(self, synthetic_dataarray, crop_geometry_bbox):
+        da = synthetic_dataarray
+        geom = list(crop_geometry_bbox)
+        with pytest.raises(ValueError):
+            _ = crop(da, geom)