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
+        - 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,31 @@
 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 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)