diff --git a/pyproject.toml b/pyproject.toml index c9f25ff..20b6abc 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -12,6 +12,7 @@ dependencies = [ "numpy", "distributed", "zarr", + "shapely", ] description = "Xarray extension for Synthetic Aperture Radar (SAR) data" readme = "README.md" diff --git a/sarxarray/utils.py b/sarxarray/utils.py index 90c311a..8008d8c 100644 --- a/sarxarray/utils.py +++ b/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): diff --git a/tests/test_utils.py b/tests/test_utils.py index e557c2f..c0e07c9 100644 --- a/tests/test_utils.py +++ b/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)