polygon_clip.py

"""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

import numpy as np
import xarray as xr

try:
    import dask.array as da
except ImportError:
    da = None

from xrspatial.utils import (_dask_task_name_kwargs, _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, all_touched=False):
    """Slice the raster to the bounding box of the geometry.

    Returns the sliced DataArray and the geometry pairs (unchanged).

    When *all_touched* is True the bounding-box comparison is expanded by
    half a pixel on every side so that pixels whose cells overlap the
    geometry boundary are not prematurely excluded.
    """
    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

    # When all_touched is set, expand the bbox by half a pixel so that
    # pixels whose cells overlap the geometry survive the crop.
    if all_touched:
        if len(x_coords) > 1:
            half_px_x = abs(float(x_coords[1] - x_coords[0])) / 2.0
        else:
            half_px_x = 0.0
        if len(y_coords) > 1:
            half_py_y = abs(float(y_coords[1] - y_coords[0])) / 2.0
        else:
            half_py_y = 0.0
        minx -= half_px_x
        maxx += half_px_x
        miny -= half_py_y
        maxy += half_py_y

    y_mask = (y_coords >= miny) & (y_coords <= maxy)
    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 attributes as *raster*.  The dtype is
        the input dtype promoted against *nodata* (e.g. an integer raster
        clipped with the default ``np.nan`` is returned as float), so the
        result is identical across the numpy, cupy, dask, and dask+cupy
        backends.

    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, all_touched=all_touched)

    # Build a binary mask via rasterize, aligned to the (possibly cropped)
    # raster grid.  Propagate the raster's chunk structure so the mask is
    # built lazily for dask backends instead of materializing a full numpy
    # array.
    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

    if has_dask_array() and isinstance(raster.data, da.Array):
        rc, cc = raster.data.chunks[-2], raster.data.chunks[-1]
        # Use the largest chunk per axis rather than the first. After a
        # crop, _crop_to_bbox slices the dask array and the leading chunk
        # is often a tiny partial edge chunk; building the rasterize mask
        # at that size fragments it into many narrow chunks and explodes
        # the task graph that xarray.where then has to align (#3191).
        kw.setdefault('chunks', (max(rc), max(cc)))
        if has_cuda_and_cupy() and is_dask_cupy(raster):
            # Respect a legacy ``use_cuda`` passed via rasterize_kw --
            # defaulting ``gpu`` as well would make rasterize() see both
            # names and raise.
            if 'use_cuda' not in kw:
                kw.setdefault('gpu', True)

    mask = rasterize(geom_pairs, **kw)

    # Apply the mask.  Keep it lazy for dask backends to avoid
    # materializing the full mask into RAM (which would OOM for 30TB
    # inputs).  For non-dask backends, compute the mask eagerly.
    mask_data = mask.data

    # Determine the output dtype the same way ``xarray.where`` does: promote
    # the raster dtype against the nodata value.  This keeps the cupy and
    # dask+cupy paths (which assign nodata into a raw copy) consistent with
    # the numpy / dask+numpy paths.  Without it, assigning ``np.nan`` into an
    # integer array raises ``cannot convert float NaN to integer`` only on the
    # GPU backends while the CPU backends silently upcast to float.
    out_dtype = np.result_type(raster.dtype, nodata)

    if has_dask_array() and isinstance(raster.data, da.Array):
        # Dask path: keep mask lazy -- no .compute()
        if isinstance(mask_data, da.Array):
            cond = mask_data == 1
        else:
            # Mask came back non-dask despite dask input (shouldn't happen,
            # but handle gracefully)
            cond = da.from_array(
                np.asarray(mask_data == 1) if not is_cupy_array(mask_data)
                else mask_data.get() == 1,
                chunks=raster.data.chunks[-2:],
            )

        if has_cuda_and_cupy() and is_dask_cupy(raster):
            # dask+cupy: use map_blocks with both raster and condition.
            # The mask is built with a uniform chunk size taken from the
            # raster's first chunk, so its block layout can differ from a
            # raster with non-uniform chunks.  map_blocks pairs blocks
            # positionally, so align the condition to the raster's chunks
            # first or it raises (or stamps the mask onto the wrong cells).
            cond = cond.rechunk(raster.data.chunks[-2:])

            def _apply_mask(raster_block, cond_block):
                out = raster_block.astype(out_dtype, copy=True)
                out[~cond_block.astype(bool)] = nodata
                return out

            out = da.map_blocks(
                _apply_mask, raster.data, cond,
                dtype=out_dtype,
                **_dask_task_name_kwargs('xrspatial.clip_polygon'),
            )
            result = xr.DataArray(out, dims=raster.dims, coords=raster.coords)
        else:
            # dask+numpy: xarray.where handles lazy condition natively
            result = raster.where(cond, other=nodata)
    elif has_cuda_and_cupy() and is_cupy_array(raster.data):
        # Pure CuPy: operate on raw arrays to avoid xarray/cupy
        # incompatibility in DataArray.where().
        import cupy
        if is_cupy_array(mask_data):
            cond_cp = mask_data == 1
        else:
            cond_cp = cupy.asarray(np.asarray(mask_data) == 1)
        out = raster.data.astype(out_dtype, copy=True)
        out[~cond_cp] = nodata
        result = xr.DataArray(out, dims=raster.dims, coords=raster.coords)
    else:
        # Pure numpy
        cond = np.asarray(mask_data) == 1
        result = raster.where(cond, other=nodata)

    result.attrs = raster.attrs
    result.name = name if name is not None else raster.name
    return result