diff --git a/examples/gallery/images/grdpolygon.py b/examples/gallery/images/grdpolygon.py new file mode 100644 index 00000000000..2e2eaefb321 --- /dev/null +++ b/examples/gallery/images/grdpolygon.py @@ -0,0 +1,80 @@ +""" +Clipping grid to a polygon +========================== + +The :func:`pygmt.grdcut` function allows you to extract a subregion from a +grid. In this example we use a :class:`geopandas.GeoDataFrame` +to crop the grid to the polygon's inner or outer region, invert it, or do both. +""" + +# %% +import geopandas as gpd +import pygmt +from shapely.geometry import Polygon + +fig = pygmt.Figure() + +# Define region of interest around Iceland +region = [-28, -10, 62, 68] # xmin, xmax, ymin, ymax + +# Load sample grid (3 arc-minutes global relief) in target area +grid = pygmt.datasets.load_earth_relief(resolution="03m", region=region) + +# Create a more polygon (irregular shape) around a smaller ROI +poly = Polygon( + [ + (-26, 63), + (-23, 63.5), + (-20, 64), + (-18, 65), + (-19, 66), + (-22, 66.5), + (-25, 66), + (-27, 65), + (-26, 63), + ] +) +gdf = gpd.GeoDataFrame({"geometry": [poly]}, crs="OGC:CRS84") + +# Original grid +fig.basemap( + region=region, + projection="M12c", + frame=["WSne+toriginal grid", "xa5f1", "ya2f1"], +) +fig.grdimage(grid=grid, cmap="oleron") + +# Cropped grid +fig.shift_origin(xshift="w+0.5c") +cropped_grid = pygmt.grdcut(grid=grid, polygon=gdf, crop=True) +fig.basemap( + region=region, + projection="M12c", + frame=["WSne+tcropped", "xa5f1", "ya2f1"], +) +fig.grdimage(grid=cropped_grid, cmap="oleron") + +# Inverted grid +fig.shift_origin(xshift="w+0.5c") +inverted_grid = pygmt.grdcut(grid=grid, polygon=gdf, invert=True) +fig.basemap( + region=region, + projection="M12c", + frame=["WSne+tinverted", "xa5f1", "ya2f1"], +) +fig.grdimage(grid=inverted_grid, cmap="oleron") + +# Cropped + inverted grid +fig.shift_origin(xshift="w+0.5c") +cropped_inverted_grid = pygmt.grdcut(grid=grid, polygon=gdf, crop=True, invert=True) +fig.basemap( + region=region, + projection="M12c", + frame=["WSne+tcropped and inverted", "xa5f1", "ya2f1"], +) +fig.grdimage(grid=cropped_inverted_grid, cmap="oleron") + +# Shared colorbar +fig.colorbar(frame=["x+lElevation", "y+lm"], position="JMR+o0.5c/0c+w8c") + +fig.show() diff --git a/pygmt/src/grdcut.py b/pygmt/src/grdcut.py index cf198d9dc74..aea856bfabf 100644 --- a/pygmt/src/grdcut.py +++ b/pygmt/src/grdcut.py @@ -3,6 +3,7 @@ """ from collections.abc import Sequence +from contextlib import ExitStack from typing import Literal import xarray as xr @@ -14,6 +15,7 @@ build_arg_list, data_kind, fmt_docstring, + tempfile_from_geojson, use_alias, ) @@ -21,7 +23,7 @@ @fmt_docstring -@use_alias(N="extend", S="circ_subregion", Z="z_subregion", f="coltypes") +@use_alias(N="extend", S="circ_subregion", Z="z_subregion", f="coltypes", F="polygon") def grdcut( grid: PathLike | xr.DataArray, kind: Literal["grid", "image"] = "grid", @@ -30,6 +32,8 @@ def grdcut( region: Sequence[float | str] | str | None = None, verbose: Literal["quiet", "error", "warning", "timing", "info", "compat", "debug"] | bool = False, + crop: bool = False, + invert: bool = False, **kwargs, ) -> xr.DataArray | None: r""" @@ -86,6 +90,17 @@ def grdcut( NaNs, append **+N** to strip off such columns before (optionally) considering the range of the core subset for further reduction of the area. + polygon : str, geopandas.GeoDataFrame, or shapely.geometry + Extract a subregion using a polygon. Can be either: + - A GMT ASCII polygon file (`.gmt`) + - A geopandas.GeoDataFrame + - A shapely.geometry.Polygon or MultiPolygon + The polygon can have holes or multiple rings. + Optional modifiers: + - crop : bool + If True, crop the output grid region to the bounding box of the polygon. + - invert : bool + If True, invert the selection, setting all nodes inside the polygon to NaN. {verbose} {coltypes} @@ -130,12 +145,33 @@ def grdcut( ) aliasdict.merge(kwargs) - with Session() as lib: + with Session() as lib, ExitStack() as stack: with ( lib.virtualfile_in(check_kind="raster", data=grid) as vingrd, lib.virtualfile_out(kind=outkind, fname=outgrid) as voutgrd, ): aliasdict["G"] = voutgrd + + if "F" in kwargs and kwargs["F"] is not None: + polygon_input = kwargs["F"] + modifiers = ("+c" * crop) + ("+i" * invert) + + # if file path provided + if isinstance(polygon_input, PathLike): + aliasdict["F"] = str(kwargs["F"]) + modifiers + # assuming its geojson + elif hasattr(polygon_input, "__geo_interface__"): + tmpfile = stack.enter_context(tempfile_from_geojson(polygon_input)) + aliasdict["F"] = tmpfile + modifiers + else: + polygon_type = type(polygon_input).__name__ + err_msg = "Invalid polygon type" + raise GMTValueError( + err_msg, + polygon_type, + "Must be a PathLike, GeoDataFrame, or Shapely geometry.", + ) + lib.call_module( module="grdcut", args=build_arg_list(aliasdict, infile=vingrd) ) diff --git a/pygmt/tests/test_grdcut.py b/pygmt/tests/test_grdcut.py index 1cd54caf4b6..3d9632e3ed6 100644 --- a/pygmt/tests/test_grdcut.py +++ b/pygmt/tests/test_grdcut.py @@ -2,6 +2,7 @@ Test pygmt.grdcut. """ +import geopandas as gpd import numpy as np import pytest import xarray as xr @@ -9,6 +10,23 @@ from pygmt.exceptions import GMTTypeError, GMTValueError from pygmt.helpers import GMTTempFile from pygmt.helpers.testing import load_static_earth_relief +from shapely.geometry import MultiPolygon, Polygon + + +@pytest.fixture(scope="module", name="polygon") +def fixture_polygon(): + """ + Provide a reusable polygon geometry for grdcut tests. + """ + return Polygon( + [ + (-52.5, -19.5), + (-51.5, -19.0), + (-50.5, -18.5), + (-51.5, -18.0), + (-52.5, -19.5), + ] + ) @pytest.fixture(scope="module", name="grid") @@ -78,3 +96,101 @@ def test_grdcut_invalid_kind(grid, region): """ with pytest.raises(GMTValueError): grdcut(grid, kind="invalid", region=region) + + +def test_grdcut_with_shapely_polygon(grid, polygon): + """ + Grdcut should accept a shapely.geometry.Polygon as polygon input + and produce a grid of expected size with some valid data. + """ + outgrid = grdcut(grid=grid, polygon=polygon) + assert outgrid is not None + + # Check size + assert outgrid.size > 0 + + # There should be some non-NaN pixels inside the polygon + assert np.any(~np.isnan(outgrid.values)) + + # Optionally, corners outside the polygon should be NaN + corner_values = [ + outgrid.values[0, 0], + outgrid.values[0, -1], + outgrid.values[-1, 0], + outgrid.values[-1, -1], + ] + assert any(np.isnan(val) for val in corner_values) + + +def test_grdcut_with_geodataframe_polygon(grid, polygon): + """ + Grdcut should accept a geopandas.GeoDataFrame as polygon input + and produce a grid of expected size with some valid data. + """ + gdf = gpd.GeoDataFrame({"geometry": [polygon]}, crs="OGC:CRS84") + outgrid = grdcut(grid=grid, polygon=gdf) + assert outgrid is not None + + # Check size + assert outgrid.size > 0 + + # Some pixels should remain non-NaN + assert np.any(~np.isnan(outgrid.values)) + + +def test_grdcut_with_polygon_file(grid, tmp_path): + """ + Grdcut should accept a GMT ASCII polygon file as polygon input. + """ + gmtfile = tmp_path / "poly.gmt" + gmtfile.write_text( + ">\n-52.5 -19.5\n-51.5 -19.0\n-50.5 -18.5\n-51.5 -18.0\n-52.5 -19.5\n" + ) + + outgrid = grdcut(grid=grid, polygon=gmtfile) + assert outgrid is not None + assert outgrid.size > 0 + + +@pytest.mark.parametrize( + ("crop", "invert"), [(True, False), (False, True), (True, True)] +) +def test_grdcut_polygon_with_crop_and_invert(grid, polygon, crop, invert): + """ + Grdcut should support crop (+c) and invert (+i) modifiers with polygon input. + """ + gdf = gpd.GeoDataFrame({"geometry": [polygon]}, crs="OGC:CRS84") + outgrid = grdcut(grid=grid, polygon=gdf, crop=crop, invert=invert) + assert outgrid is not None + assert outgrid.size > 0 + + if invert: + assert np.isnan(outgrid.data).any() + + assert np.isfinite(outgrid.data).any() + + assert np.count_nonzero(np.isfinite(outgrid.data)) < grid.size + + +def test_grdcut_with_multipolygon(grid): + """ + Grdcut should accept a shapely.geometry.MultiPolygon as polygon input. + """ + # Create two simple polygons and combine into a MultiPolygon + poly1 = Polygon([(-52.5, -19.5), (-51.5, -19.0), (-50.5, -18.5), (-52.5, -19.5)]) + poly2 = Polygon([(-51.0, -18.5), (-50.5, -18.0), (-50.0, -18.5), (-51.0, -18.5)]) + multipoly = MultiPolygon([poly1, poly2]) + + outgrid = grdcut(grid=grid, polygon=multipoly) + assert outgrid is not None + # Ensure some grid pixels exist + assert outgrid.size > 0 + assert np.any(np.isnan(outgrid.values)) + + +def test_grdcut_polygon_invalid_input(grid): + """ + Grdcut should raise an error for invalid polygon input type. + """ + with pytest.raises(GMTValueError): + grdcut(grid=grid, polygon=12345)