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,39 @@ def _compute_coherence(numerator, denominator):
     return coherence
 
 
+def crop(data: xr.Dataset | xr.DataArray, geom: sg.Polygon) -> 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
+        shapely.geometry.Polygon in radar coordinates of the area that should be
+        kept, in [azimuth, 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 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)
+    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,17 @@ 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],
+        [602, 1405]
+    ])
 
 # this class tests multi_look with dataarray. For testing with dataset, see
 # test_stack.py
@@ -276,3 +289,20 @@ 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(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_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)