Support polygon-based subsetting in grdcut -F with GeoDataFrame

examples/gallery/images/grdpolygon.py

@@ -0,0 +1,80 @@
+"""
+Clipping grid to a complex polygon
+==================================
+
+The :func:`pygmt.grdcut` function allows you to extract a subregion from a
+grid. In this example we use a complex polygon (GeoDataFrame or GMT ASCII file)
+to crop the grid, 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]
+
+# Load sample grid (3 arc-minutes global relief) in target area
+grid = pygmt.datasets.load_earth_relief(resolution="03m", region=region)
+
+# Create a more complex polygon (irregular shape) around a smaller ROI
+complex_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": [complex_poly]}, crs="EPSG:4326")
+
+# 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()

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,14 +15,15 @@
     build_arg_list,
     data_kind,
     fmt_docstring,
+    tempfile_from_geojson,
     use_alias,
 )
 
 __doctest_skip__ = ["grdcut"]
 
 
 @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 (must have CRS EPSG:4326)
+        - 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,23 @@ 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 not isinstance(polygon_input, (str, bytes)):
+                    tmpfile = stack.enter_context(tempfile_from_geojson(polygon_input))
+                    aliasdict["F"] = tmpfile + modifiers
+                else:
+                    aliasdict["F"] = str(polygon_input) + modifiers
+
             lib.call_module(
                 module="grdcut", args=build_arg_list(aliasdict, infile=vingrd)
             )