Support polygon-based subsetting in grdcut -F with GeoDataFrame

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()

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

pygmt/tests/test_grdcut.py

@@ -2,13 +2,26 @@
 Test pygmt.grdcut.
 """
 
+import geopandas as gpd
 import numpy as np
 import pytest
 import xarray as xr
 from pygmt import grdcut
 from pygmt.exceptions import GMTTypeError, GMTValueError
 from pygmt.helpers import GMTTempFile
 from pygmt.helpers.testing import load_static_earth_relief
+from shapely.geometry import Polygon
+
+# Reusable polygon geometry
+poly = 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 +91,58 @@ def test_grdcut_invalid_kind(grid, region):
     """
     with pytest.raises(GMTValueError):
         grdcut(grid, kind="invalid", region=region)
+
+
+def test_grdcut_with_shapely_polygon(grid):
+    """
+    Grdcut should accept a shapely.geometry.Polygon as polygon input.
+    """
+    outgrid = grdcut(grid=grid, polygon=poly)
+    assert outgrid is not None
+    assert outgrid.rio.crs is None or str(outgrid.rio.crs) == ""  # CRS not enforced
+    assert outgrid.size > 0
+
+
+def test_grdcut_with_geodataframe_polygon(grid):
+    """
+    Grdcut should accept a geopandas.GeoDataFrame as polygon input.
+    """
+    gdf = gpd.GeoDataFrame({"geometry": [poly]}, crs="OGC:CRS84")
+    outgrid = grdcut(grid=grid, polygon=gdf)
+    assert outgrid is not None
+    assert outgrid.size > 0
+
+
+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, crop, invert):
+    """
+    Grdcut should support crop (+c) and invert (+i) modifiers with polygon input.
+    """
+    gdf = gpd.GeoDataFrame({"geometry": [poly]}, crs="OGC:CRS84")
+    outgrid = grdcut(grid=grid, polygon=gdf, crop=crop, invert=invert)
+    assert outgrid is not None
+    assert outgrid.size > 0
+
+
+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)