64 Native cropping in radar coordinates

pyproject.toml

@@ -12,6 +12,7 @@ dependencies = [
     "numpy",
     "distributed",
     "zarr",
+    "shapely",
 ]
 description = "Xarray extension for Synthetic Aperture Radar (SAR) data"
 readme = "README.md"

sarxarray/utils.py

@@ -1,4 +1,5 @@
 import numpy as np
+import shapely.geometry as sg
 import xarray as xr
 from dask.delayed import Delayed, delayed
 
@@ -154,6 +155,59 @@ def _compute_coherence(numerator, denominator):
     return coherence
 
 
+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 | tuple
+        shapely.geometry.Polygon in radar coordinates of the area that should be
+        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
+    -------
+    xr.Dataset | xr.DataArray
+        The dataset or data array cropped to the area of interest
+
+    Raises
+    ------
+    ValueError
+        - 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.")
+
+    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)
+    )
+
+    return data
+
+
 def _validate_multi_look_inputs(data, window_size, method, statistics):
     # check if data is xarray
     if not isinstance(data, xr.Dataset | xr.DataArray):

tests/test_utils.py

@@ -2,11 +2,13 @@
 import pytest
 import xarray as xr
 from dask.delayed import Delayed
+from shapely.geometry import Polygon
 
 from sarxarray.utils import (
     _get_chunks,
     _validate_multi_look_inputs,
     complex_coherence,
+    crop,
     multi_look,
 )
 
@@ -27,6 +29,22 @@ def synthetic_dataarray():
         },
     )
 
+# Create a polygon for cropping
+@pytest.fixture
+def crop_geometry():
+    return Polygon([
+        [601.9, 1405],
+        [605, 1407],
+        [606, 1408],
+        [606, 1409],
+        [603, 1409],
+        [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
@@ -276,3 +294,62 @@ def test_complex_coherence_bad_args(
             complex_coherence(reference, other1, window_size=(2, 2), compute=True)
         with pytest.raises(ValueError):
             complex_coherence(reference, other2, window_size=(2, 2), compute=True)
+
+
+class TestUtilsCrop:
+    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)