diff --git a/README.md b/README.md index c04b640f..62e6b734 100644 --- a/README.md +++ b/README.md @@ -460,6 +460,7 @@ Same-CRS tiles skip reprojection entirely and are placed by direct coordinate al | Name | Description | Source | NumPy xr.DataArray | Dask xr.DataArray | CuPy GPU xr.DataArray | Dask GPU xr.DataArray | |:----------:|:------------|:------:|:----------------------:|:--------------------:|:-------------------:|:------:| | [Apply](xrspatial/zonal.py) | Applies a custom function to each zone in a classified raster | Standard | ✅️ | ✅️ | ✅️ | ✅️ | +| [Clip Polygon](xrspatial/polygon_clip.py) | Clips a raster to an arbitrary polygon with masking | Standard | ✅️ | ✅️ | ✅️ | ✅️ | | [Crop](xrspatial/zonal.py) | Extracts the bounding rectangle of a specific zone | Standard | ✅️ | ✅️ | ✅️ | ✅️ | | [Regions](xrspatial/zonal.py) | Identifies connected regions of non-zero cells | Standard (CCL) | ✅️ | ✅️ | ✅️ | ✅️ | | [Trim](xrspatial/zonal.py) | Removes nodata border rows and columns from a raster | Standard | ✅️ | ✅️ | ✅️ | ✅️ | diff --git a/docs/source/reference/zonal.rst b/docs/source/reference/zonal.rst index 7d20706d..13c1200c 100644 --- a/docs/source/reference/zonal.rst +++ b/docs/source/reference/zonal.rst @@ -16,6 +16,13 @@ Apply xrspatial.zonal.apply +Clip Polygon +============ +.. autosummary:: + :toctree: _autosummary + + xrspatial.polygon_clip.clip_polygon + Crop ==== .. autosummary:: diff --git a/examples/user_guide/50_Polygon_Clipping.ipynb b/examples/user_guide/50_Polygon_Clipping.ipynb new file mode 100644 index 00000000..ec54b13d --- /dev/null +++ b/examples/user_guide/50_Polygon_Clipping.ipynb @@ -0,0 +1,311 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "a1", + "metadata": {}, + "source": [ + "# Polygon Clipping\n", + "\n", + "Extracting raster data for an irregular boundary is one of the most common GIS operations. `clip_polygon()` masks pixels that fall outside an arbitrary polygon to a nodata value (NaN by default) and optionally trims the result to the polygon's bounding box.\n", + "\n", + "This is different from `crop()`, which only extracts a rectangular bounding box. When your study area is a watershed, administrative boundary, or any non-rectangular shape, `clip_polygon()` is what you want." + ] + }, + { + "cell_type": "markdown", + "id": "a2", + "metadata": {}, + "source": [ + "### What you'll see\n", + "\n", + "1. Generate a synthetic terrain raster\n", + "2. Clip to a single polygon\n", + "3. Compare crop=True vs crop=False\n", + "4. Clip with a custom nodata value\n", + "5. Use all_touched to include boundary pixels\n", + "6. Clip with a list of coordinate pairs" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a3", + "metadata": {}, + "outputs": [], + "source": [ + "import numpy as np\n", + "import xarray as xr\n", + "import matplotlib.pyplot as plt\n", + "from matplotlib.patches import Polygon as MplPolygon\n", + "from shapely.geometry import Polygon\n", + "\n", + "from xrspatial import generate_terrain\n", + "from xrspatial.polygon_clip import clip_polygon" + ] + }, + { + "cell_type": "markdown", + "id": "a4", + "metadata": {}, + "source": [ + "## Generate terrain\n", + "\n", + "A 200x300 synthetic elevation raster gives us something realistic to clip." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a5", + "metadata": {}, + "outputs": [], + "source": [ + "template = xr.DataArray(np.zeros((200, 300)))\n", + "terrain = generate_terrain(template, x_range=(0, 30), y_range=(0, 20))\n", + "\n", + "fig, ax = plt.subplots(figsize=(10, 6))\n", + "terrain.plot(ax=ax, cmap='terrain', add_colorbar=True)\n", + "ax.set_title('Synthetic terrain')\n", + "ax.set_aspect('equal')\n", + "plt.tight_layout()\n", + "plt.show()\n", + "\n", + "print(f'Shape: {terrain.shape}')" + ] + }, + { + "cell_type": "markdown", + "id": "a6", + "metadata": {}, + "source": [ + "## Define a clipping polygon\n", + "\n", + "An irregular polygon representing a study area boundary. We'll draw it on top of the terrain to see what we're about to extract." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a7", + "metadata": {}, + "outputs": [], + "source": [ + "study_area = Polygon([\n", + " (5, 4), (8, 2), (14, 3), (20, 5),\n", + " (22, 10), (18, 16), (12, 17),\n", + " (6, 14), (3, 9),\n", + "])\n", + "\n", + "fig, ax = plt.subplots(figsize=(10, 6))\n", + "terrain.plot(ax=ax, cmap='terrain', add_colorbar=True)\n", + "patch = MplPolygon(\n", + " list(study_area.exterior.coords),\n", + " closed=True, fill=False, edgecolor='red', linewidth=2,\n", + ")\n", + "ax.add_patch(patch)\n", + "ax.set_title('Terrain with study area boundary')\n", + "ax.set_aspect('equal')\n", + "plt.tight_layout()\n", + "plt.show()" + ] + }, + { + "cell_type": "markdown", + "id": "a8", + "metadata": {}, + "source": [ + "## Basic clip\n", + "\n", + "With `crop=True` (the default), the result is trimmed to the polygon's bounding box. Pixels outside the polygon are set to NaN." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a9", + "metadata": {}, + "outputs": [], + "source": [ + "clipped = clip_polygon(terrain, study_area)\n", + "\n", + "fig, axes = plt.subplots(1, 2, figsize=(14, 5))\n", + "\n", + "terrain.plot(ax=axes[0], cmap='terrain', add_colorbar=False)\n", + "axes[0].set_title('Original')\n", + "axes[0].set_aspect('equal')\n", + "\n", + "clipped.plot(ax=axes[1], cmap='terrain', add_colorbar=False)\n", + "axes[1].set_title('Clipped (crop=True)')\n", + "axes[1].set_aspect('equal')\n", + "\n", + "plt.tight_layout()\n", + "plt.show()\n", + "\n", + "print(f'Original shape: {terrain.shape}')\n", + "print(f'Clipped shape: {clipped.shape}')" + ] + }, + { + "cell_type": "markdown", + "id": "a10", + "metadata": {}, + "source": [ + "## crop=False: keep the original extent\n", + "\n", + "When `crop=False`, the output has the same shape as the input. Pixels outside the polygon are still masked to NaN, but the spatial extent is unchanged. This is useful when you need the result to align pixel-for-pixel with other rasters." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a11", + "metadata": {}, + "outputs": [], + "source": [ + "clipped_full = clip_polygon(terrain, study_area, crop=False)\n", + "\n", + "fig, ax = plt.subplots(figsize=(10, 6))\n", + "clipped_full.plot(ax=ax, cmap='terrain', add_colorbar=True)\n", + "ax.set_title('Clipped (crop=False) — same extent, masked outside polygon')\n", + "ax.set_aspect('equal')\n", + "plt.tight_layout()\n", + "plt.show()\n", + "\n", + "print(f'Shape matches original: {clipped_full.shape == terrain.shape}')" + ] + }, + { + "cell_type": "markdown", + "id": "a12", + "metadata": {}, + "source": [ + "## Custom nodata value\n", + "\n", + "By default, masked pixels get NaN. You can set a different value with the `nodata` parameter. This is helpful for integer rasters or when downstream tools expect a specific sentinel value." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a13", + "metadata": {}, + "outputs": [], + "source": [ + "clipped_nd = clip_polygon(terrain, study_area, nodata=-9999, crop=False)\n", + "\n", + "# Count cells with each status\n", + "vals = clipped_nd.values\n", + "n_nodata = np.sum(vals == -9999)\n", + "n_valid = np.sum(vals != -9999)\n", + "print(f'Valid pixels: {n_valid}')\n", + "print(f'Nodata pixels: {n_nodata}')" + ] + }, + { + "cell_type": "markdown", + "id": "a14", + "metadata": {}, + "source": [ + "## all_touched: boundary pixel inclusion\n", + "\n", + "By default, only pixels whose centre falls inside the polygon are kept. With `all_touched=True`, any pixel that the polygon boundary touches is also included. This gives a slightly larger footprint." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a15", + "metadata": {}, + "outputs": [], + "source": [ + "clipped_default = clip_polygon(terrain, study_area, crop=False)\n", + "clipped_touched = clip_polygon(terrain, study_area, crop=False, all_touched=True)\n", + "\n", + "n_default = np.count_nonzero(np.isfinite(clipped_default.values))\n", + "n_touched = np.count_nonzero(np.isfinite(clipped_touched.values))\n", + "\n", + "fig, axes = plt.subplots(1, 2, figsize=(14, 5))\n", + "\n", + "clipped_default.plot(ax=axes[0], cmap='terrain', add_colorbar=False)\n", + "axes[0].set_title(f'all_touched=False ({n_default} pixels)')\n", + "axes[0].set_aspect('equal')\n", + "\n", + "clipped_touched.plot(ax=axes[1], cmap='terrain', add_colorbar=False)\n", + "axes[1].set_title(f'all_touched=True ({n_touched} pixels)')\n", + "axes[1].set_aspect('equal')\n", + "\n", + "plt.tight_layout()\n", + "plt.show()\n", + "\n", + "print(f'all_touched adds {n_touched - n_default} boundary pixels')" + ] + }, + { + "cell_type": "markdown", + "id": "a16", + "metadata": {}, + "source": [ + "## Coordinate array input\n", + "\n", + "If you don't have shapely installed in your workflow, you can pass the polygon as a list of (x, y) coordinate pairs. `clip_polygon()` builds the polygon internally." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a17", + "metadata": {}, + "outputs": [], + "source": [ + "coords = [(5, 4), (8, 2), (14, 3), (20, 5),\n", + " (22, 10), (18, 16), (12, 17), (6, 14), (3, 9)]\n", + "\n", + "clipped_coords = clip_polygon(terrain, coords)\n", + "\n", + "fig, ax = plt.subplots(figsize=(10, 6))\n", + "clipped_coords.plot(ax=ax, cmap='terrain', add_colorbar=True)\n", + "ax.set_title('Clipped from coordinate list')\n", + "ax.set_aspect('equal')\n", + "plt.tight_layout()\n", + "plt.show()" + ] + }, + { + "cell_type": "markdown", + "id": "a18", + "metadata": {}, + "source": [ + "## Accessor syntax\n", + "\n", + "All xrspatial functions are also available through the `.xrs` accessor on DataArrays." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a19", + "metadata": {}, + "outputs": [], + "source": [ + "import xrspatial # registers the .xrs accessor\n", + "\n", + "clipped_acc = terrain.xrs.clip_polygon(study_area)\n", + "print(f'Accessor result shape: {clipped_acc.shape}')" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python", + "version": "3.14.2" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/xrspatial/__init__.py b/xrspatial/__init__.py index 643f3ac5..6c02280c 100644 --- a/xrspatial/__init__.py +++ b/xrspatial/__init__.py @@ -120,6 +120,7 @@ from xrspatial.terrain_metrics import tri # noqa from xrspatial.hydro import twi # noqa: unified wrapper from xrspatial.hydro import twi_d8 # noqa +from xrspatial.polygon_clip import clip_polygon # noqa from xrspatial.polygonize import polygonize # noqa from xrspatial.viewshed import viewshed # noqa from xrspatial.visibility import cumulative_viewshed # noqa diff --git a/xrspatial/accessor.py b/xrspatial/accessor.py index 2ab83627..bd3bb35d 100644 --- a/xrspatial/accessor.py +++ b/xrspatial/accessor.py @@ -362,6 +362,10 @@ def crop(self, zones, zones_ids, **kwargs): from .zonal import crop return crop(zones, self._obj, zones_ids, **kwargs) + def clip_polygon(self, geometry, **kwargs): + from .polygon_clip import clip_polygon + return clip_polygon(self._obj, geometry, **kwargs) + def trim(self, **kwargs): from .zonal import trim return trim(self._obj, **kwargs) diff --git a/xrspatial/polygon_clip.py b/xrspatial/polygon_clip.py new file mode 100644 index 00000000..f8e2fe65 --- /dev/null +++ b/xrspatial/polygon_clip.py @@ -0,0 +1,247 @@ +"""Clip a raster to an arbitrary polygon geometry. + +Unlike :func:`~xrspatial.zonal.crop`, which extracts a rectangular bounding +box, ``clip_polygon`` masks pixels that fall outside the polygon boundary to +a configurable nodata value (default ``np.nan``). +""" + +from __future__ import annotations + +from typing import Optional, Sequence, Tuple, Union + +import numpy as np +import xarray as xr + +try: + import dask.array as da +except ImportError: + da = None + +from xrspatial.utils import ( + _validate_raster, has_cuda_and_cupy, has_dask_array, is_cupy_array, + is_dask_cupy, +) + + +def _resolve_geometry(geometry): + """Return a list of ``(shapely_geom, 1.0)`` pairs for rasterize(). + + Accepts a single shapely geometry, an iterable of shapely geometries, + a GeoDataFrame/GeoSeries, or a coordinate array ``[(x, y), ...]``. + """ + # GeoDataFrame / GeoSeries — check FIRST because GeoDataFrame has a + # ``geom_type`` property that returns a *Series*, which would break the + # scalar ``in`` check below. + try: + import geopandas as gpd + if isinstance(geometry, gpd.GeoDataFrame): + from shapely.ops import unary_union + merged = unary_union(geometry.geometry) + return [(merged, 1.0)] + if isinstance(geometry, gpd.GeoSeries): + from shapely.ops import unary_union + merged = unary_union(geometry) + return [(merged, 1.0)] + except ImportError: + pass + + # Single shapely geometry + _base = getattr(geometry, 'geom_type', None) + if isinstance(_base, str) and _base in ('Polygon', 'MultiPolygon'): + return [(geometry, 1.0)] + + # Iterable of shapely geometries or coordinate pairs + if hasattr(geometry, '__iter__'): + items = list(geometry) + if len(items) == 0: + raise ValueError("geometry is empty") + + first = items[0] + + # List of shapely geometries + ft = getattr(first, 'geom_type', None) + if isinstance(ft, str) and ft in ('Polygon', 'MultiPolygon'): + from shapely.ops import unary_union + merged = unary_union(items) + return [(merged, 1.0)] + + # Coordinate array: [(x, y), ...] — build a polygon + if (isinstance(first, (list, tuple, np.ndarray)) + and len(first) == 2): + from shapely.geometry import Polygon as ShapelyPolygon + poly = ShapelyPolygon(items) + return [(poly, 1.0)] + + raise TypeError( + f"geometry must be a shapely Polygon/MultiPolygon, an iterable of " + f"geometries, a GeoDataFrame, or a list of (x, y) coordinates. " + f"Got {type(geometry)}" + ) + + +def _crop_to_bbox(raster, geom_pairs): + """Slice the raster to the bounding box of the geometry. + + Returns the sliced DataArray and the geometry pairs (unchanged). + """ + from shapely.ops import unary_union + merged = unary_union([g for g, _ in geom_pairs]) + minx, miny, maxx, maxy = merged.bounds + + y_coords = raster.coords[raster.dims[-2]].values + x_coords = raster.coords[raster.dims[-1]].values + + # Determine coordinate ordering (ascending or descending) + y_ascending = y_coords[-1] >= y_coords[0] + x_ascending = x_coords[-1] >= x_coords[0] + + if y_ascending: + y_mask = (y_coords >= miny) & (y_coords <= maxy) + else: + y_mask = (y_coords >= miny) & (y_coords <= maxy) + + if x_ascending: + x_mask = (x_coords >= minx) & (x_coords <= maxx) + else: + x_mask = (x_coords >= minx) & (x_coords <= maxx) + + y_idx = np.where(y_mask)[0] + x_idx = np.where(x_mask)[0] + + if len(y_idx) == 0 or len(x_idx) == 0: + raise ValueError( + "Clipping geometry does not overlap the raster extent." + ) + + y_slice = slice(int(y_idx[0]), int(y_idx[-1]) + 1) + x_slice = slice(int(x_idx[0]), int(x_idx[-1]) + 1) + + return raster[..., y_slice, x_slice] + + +def clip_polygon( + raster: xr.DataArray, + geometry, + nodata: float = np.nan, + crop: bool = True, + all_touched: bool = False, + rasterize_kw: Optional[dict] = None, + name: Optional[str] = None, +) -> xr.DataArray: + """Clip a raster to an arbitrary polygon geometry. + + Pixels outside the polygon are set to *nodata*. When *crop* is True + (the default), the output is also trimmed to the polygon's bounding + box so the result is smaller than the input. + + Parameters + ---------- + raster : xr.DataArray + Input raster to clip. Must be at least 2-D with named ``y`` + and ``x`` dimensions (last two dimensions). + geometry : shapely geometry, list of geometries, GeoDataFrame, or coordinate array + The clipping polygon(s). Accepts: + + * A single ``shapely.geometry.Polygon`` or ``MultiPolygon``. + * An iterable of shapely polygon geometries (merged via + ``unary_union``). + * A ``GeoDataFrame`` or ``GeoSeries`` (merged via + ``unary_union``). + * A list of ``(x, y)`` coordinate pairs defining a single + polygon ring. + nodata : float, default np.nan + Value to assign to pixels outside the polygon. + crop : bool, default True + If True, also trim the output to the bounding box of the + polygon. This reduces memory usage for small clip regions + within large rasters. + all_touched : bool, default False + If True, all pixels touched by the polygon boundary are + included. If False (default), only pixels whose centre falls + inside the polygon are included. + rasterize_kw : dict, optional + Extra keyword arguments forwarded to :func:`~xrspatial.rasterize.rasterize` + when creating the polygon mask. + name : str, optional + Name for the output DataArray. Defaults to the input name. + + Returns + ------- + xr.DataArray + Clipped raster with the same dtype and attributes as *raster*. + + Examples + -------- + .. sourcecode:: python + + >>> from shapely.geometry import Polygon + >>> import xrspatial + >>> poly = Polygon([(1, 1), (1, 3), (3, 3), (3, 1)]) + >>> clipped = xrspatial.clip_polygon(raster, poly) + """ + _validate_raster(raster, func_name='clip_polygon', name='raster') + + # Resolve geometry into [(shapely_geom, 1.0)] pairs + geom_pairs = _resolve_geometry(geometry) + + # Optionally crop to bounding box first (reduces rasterize cost) + if crop: + raster = _crop_to_bbox(raster, geom_pairs) + + # Build a binary mask via rasterize, aligned to the (possibly cropped) + # raster grid. Always produce a plain numpy mask first, then convert + # to the raster's backend so xarray's .where() sees matching types. + from .rasterize import rasterize + + kw = dict(rasterize_kw or {}) + kw['like'] = raster + kw['fill'] = 0.0 + kw['dtype'] = np.uint8 + kw['all_touched'] = all_touched + + mask = rasterize(geom_pairs, **kw) + + # Get mask as a plain numpy array regardless of what rasterize returned + mask_np = mask.data + if has_dask_array() and isinstance(mask_np, da.Array): + mask_np = mask_np.compute() + if has_cuda_and_cupy() and is_cupy_array(mask_np): + mask_np = mask_np.get() + + # Apply the mask. For CuPy backends we operate on the raw array + # to avoid an xarray/cupy incompatibility in ``DataArray.where()``. + cond_np = mask_np == 1 + + if has_cuda_and_cupy() and is_cupy_array(raster.data): + import cupy + cond_cp = cupy.asarray(cond_np) + out = raster.data.copy() + out[~cond_cp] = nodata + result = xr.DataArray(out, dims=raster.dims, coords=raster.coords) + elif has_cuda_and_cupy() and is_dask_cupy(raster): + import cupy + cond_cp = cupy.asarray(cond_np) + + def _apply_mask(block, block_info=None): + if block_info is not None: + loc = block_info[0]['array-location'] + y_sl = slice(loc[-2][0], loc[-2][1]) + x_sl = slice(loc[-1][0], loc[-1][1]) + c = cond_cp[y_sl, x_sl] + else: + c = cond_cp + out = block.copy() + out[~c] = nodata + return out + + out = raster.data.map_blocks(_apply_mask, dtype=raster.dtype) + result = xr.DataArray(out, dims=raster.dims, coords=raster.coords) + elif has_dask_array() and isinstance(raster.data, da.Array): + cond_da = da.from_array(cond_np, chunks=raster.data.chunks[-2:]) + result = raster.where(cond_da, other=nodata) + else: + result = raster.where(cond_np, other=nodata) + + result.attrs = raster.attrs + result.name = name if name is not None else raster.name + return result diff --git a/xrspatial/tests/test_polygon_clip.py b/xrspatial/tests/test_polygon_clip.py new file mode 100644 index 00000000..e0645662 --- /dev/null +++ b/xrspatial/tests/test_polygon_clip.py @@ -0,0 +1,274 @@ +"""Tests for xrspatial.polygon_clip.clip_polygon.""" + +from __future__ import annotations + +import numpy as np +import pytest +import xarray as xr + +from shapely.geometry import Polygon, MultiPolygon, box + +from xrspatial.polygon_clip import clip_polygon +from xrspatial.tests.general_checks import ( + create_test_raster, + cuda_and_cupy_available, + dask_array_available, +) + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +def _make_raster(backend='numpy', chunks=(3, 3)): + """8x6 raster with known values and coordinates spanning [0, 2.5] x [0, 3.5].""" + data = np.arange(48, dtype=np.float64).reshape(8, 6) + return create_test_raster(data, backend=backend, chunks=chunks) + + +def _inner_polygon(): + """Polygon covering the centre of the 8x6 test raster. + + Raster y goes from 3.5 (top) to 0.0 (bottom) and x from 0.0 to 2.5. + This polygon covers roughly the inner 4x4 cell area. + """ + return Polygon([(0.6, 0.6), (0.6, 2.9), (1.9, 2.9), (1.9, 0.6)]) + + +# --------------------------------------------------------------------------- +# NumPy backend +# --------------------------------------------------------------------------- + +class TestClipPolygonNumpy: + def test_basic_mask(self): + """Pixels outside polygon are NaN, inside are preserved.""" + raster = _make_raster() + # Use a triangular polygon so corners outside the triangle + # get masked even after crop to bounding box. + poly = Polygon([(0.6, 0.6), (1.25, 2.9), (1.9, 0.6)]) + result = clip_polygon(raster, poly) + + # Result should be smaller than input (cropped to bbox) + assert result.shape[0] <= raster.shape[0] + assert result.shape[1] <= raster.shape[1] + + vals = result.values + # At least some pixels are NaN (masked corners outside triangle) + assert np.any(np.isnan(vals)) + # At least some pixels are preserved + assert np.any(np.isfinite(vals)) + + def test_no_crop(self): + """With crop=False, output shape matches input.""" + raster = _make_raster() + poly = _inner_polygon() + result = clip_polygon(raster, poly, crop=False) + + assert result.shape == raster.shape + + vals = result.values + # Pixels outside polygon are NaN + assert np.any(np.isnan(vals)) + # Pixels inside are preserved + assert np.any(np.isfinite(vals)) + + def test_nodata_value(self): + """Custom nodata value is applied outside the polygon.""" + raster = _make_raster() + poly = _inner_polygon() + result = clip_polygon(raster, poly, nodata=-9999.0, crop=False) + + vals = result.values + # No NaNs (nodata is -9999 instead) + assert not np.any(np.isnan(vals)) + # Some cells have the nodata sentinel + assert np.any(vals == -9999.0) + + def test_preserves_attrs(self): + """Output DataArray keeps the input's attributes.""" + raster = _make_raster() + poly = _inner_polygon() + result = clip_polygon(raster, poly) + assert result.attrs == raster.attrs + + def test_custom_name(self): + """Output name can be overridden.""" + raster = _make_raster() + poly = _inner_polygon() + result = clip_polygon(raster, poly, name='clipped_1144') + assert result.name == 'clipped_1144' + + def test_coordinate_array_input(self): + """Geometry given as a list of (x, y) coordinate pairs.""" + raster = _make_raster() + coords = [(0.6, 0.6), (0.6, 2.9), (1.9, 2.9), (1.9, 0.6)] + result = clip_polygon(raster, coords) + assert np.any(np.isfinite(result.values)) + + def test_multipolygon(self): + """MultiPolygon geometry is accepted.""" + raster = _make_raster() + p1 = Polygon([(0.0, 0.0), (0.0, 1.0), (1.0, 1.0), (1.0, 0.0)]) + p2 = Polygon([(1.5, 1.5), (1.5, 2.5), (2.5, 2.5), (2.5, 1.5)]) + mp = MultiPolygon([p1, p2]) + result = clip_polygon(raster, mp, crop=False) + assert result.shape == raster.shape + + def test_list_of_polygons(self): + """List of shapely polygons is merged and applied.""" + raster = _make_raster() + polys = [ + Polygon([(0.0, 0.0), (0.0, 1.0), (1.0, 1.0), (1.0, 0.0)]), + Polygon([(1.5, 1.5), (1.5, 2.5), (2.5, 2.5), (2.5, 1.5)]), + ] + result = clip_polygon(raster, polys, crop=False) + assert result.shape == raster.shape + + def test_nan_in_input_preserved(self): + """NaN values in the input stay NaN in the output.""" + data = np.arange(48, dtype=np.float64).reshape(8, 6) + data[3, 3] = np.nan + raster = create_test_raster(data) + poly = box(0, 0, 2.5, 3.5) # covers entire raster + result = clip_polygon(raster, poly, crop=False) + assert np.isnan(result.values[3, 3]) + + def test_non_overlapping_raises(self): + """Polygon completely outside the raster raises ValueError.""" + raster = _make_raster() + poly = Polygon([(100, 100), (100, 200), (200, 200), (200, 100)]) + with pytest.raises(ValueError, match="does not overlap"): + clip_polygon(raster, poly) + + def test_all_touched(self): + """all_touched=True includes boundary pixels.""" + raster = _make_raster() + poly = _inner_polygon() + result_default = clip_polygon(raster, poly, crop=False) + result_touched = clip_polygon(raster, poly, crop=False, + all_touched=True) + # all_touched should include at least as many non-NaN pixels + n_default = np.count_nonzero(np.isfinite(result_default.values)) + n_touched = np.count_nonzero(np.isfinite(result_touched.values)) + assert n_touched >= n_default + + def test_single_cell_raster(self): + """1x1 raster with polygon that covers it.""" + data = np.array([[42.0]]) + raster = create_test_raster(data, attrs={'res': (1.0, 1.0)}) + poly = box(-1, -1, 1, 1) + result = clip_polygon(raster, poly) + assert result.values[0, 0] == 42.0 + + def test_empty_geometry_raises(self): + """Empty geometry list raises.""" + raster = _make_raster() + with pytest.raises(ValueError, match="empty"): + clip_polygon(raster, []) + + +# --------------------------------------------------------------------------- +# Dask + NumPy backend +# --------------------------------------------------------------------------- + +@dask_array_available +class TestClipPolygonDask: + def test_matches_numpy(self): + """Dask result matches NumPy result.""" + np_raster = _make_raster(backend='numpy') + dk_raster = _make_raster(backend='dask+numpy', chunks=(4, 3)) + poly = _inner_polygon() + + np_result = clip_polygon(np_raster, poly, crop=False) + dk_result = clip_polygon(dk_raster, poly, crop=False) + + np.testing.assert_allclose( + dk_result.values, np_result.values, equal_nan=True + ) + + def test_lazy(self): + """Output remains a dask array (not eagerly computed).""" + import dask.array as da + raster = _make_raster(backend='dask+numpy') + poly = _inner_polygon() + result = clip_polygon(raster, poly) + assert isinstance(result.data, da.Array) + + def test_crop_matches_numpy(self): + """Cropped dask result matches cropped NumPy result.""" + np_raster = _make_raster(backend='numpy') + dk_raster = _make_raster(backend='dask+numpy', chunks=(4, 3)) + poly = _inner_polygon() + + np_result = clip_polygon(np_raster, poly, crop=True) + dk_result = clip_polygon(dk_raster, poly, crop=True) + + np.testing.assert_allclose( + dk_result.values, np_result.values, equal_nan=True + ) + + +# --------------------------------------------------------------------------- +# CuPy backend +# --------------------------------------------------------------------------- + +@cuda_and_cupy_available +class TestClipPolygonCuPy: + def test_matches_numpy(self): + """CuPy result matches NumPy result.""" + np_raster = _make_raster(backend='numpy') + cp_raster = _make_raster(backend='cupy') + poly = _inner_polygon() + + np_result = clip_polygon(np_raster, poly, crop=False) + cp_result = clip_polygon(cp_raster, poly, crop=False) + + np.testing.assert_allclose( + cp_result.data.get(), np_result.values, equal_nan=True + ) + + +# --------------------------------------------------------------------------- +# Dask + CuPy backend +# --------------------------------------------------------------------------- + +@cuda_and_cupy_available +@dask_array_available +class TestClipPolygonDaskCuPy: + def test_matches_numpy(self): + """Dask+CuPy result matches NumPy result.""" + np_raster = _make_raster(backend='numpy') + dkcp_raster = _make_raster(backend='dask+cupy', chunks=(4, 3)) + poly = _inner_polygon() + + np_result = clip_polygon(np_raster, poly, crop=False) + dkcp_result = clip_polygon(dkcp_raster, poly, crop=False) + + np.testing.assert_allclose( + dkcp_result.data.compute().get(), np_result.values, equal_nan=True + ) + + +# --------------------------------------------------------------------------- +# GeoDataFrame input +# --------------------------------------------------------------------------- + +class TestClipPolygonGeoDataFrame: + def test_geodataframe_input(self): + """GeoDataFrame geometry is accepted.""" + geopandas = pytest.importorskip("geopandas") + raster = _make_raster() + poly = _inner_polygon() + gdf = geopandas.GeoDataFrame(geometry=[poly]) + result = clip_polygon(raster, gdf, crop=False) + assert result.shape == raster.shape + assert np.any(np.isfinite(result.values)) + + def test_geoseries_input(self): + """GeoSeries geometry is accepted.""" + geopandas = pytest.importorskip("geopandas") + raster = _make_raster() + poly = _inner_polygon() + gs = geopandas.GeoSeries([poly]) + result = clip_polygon(raster, gs, crop=False) + assert result.shape == raster.shape