pygmt.grdmask: Create mask grid from polygons or point coverage(required arguments)

doc/api/index.rst

@@ -126,6 +126,7 @@ Operations on tabular data
     blockmedian
     blockmode
     filter1d
+    grdmask
     nearneighbor
     project
     select

pygmt/__init__.py

@@ -45,6 +45,7 @@
     grdhisteq,
     grdinfo,
     grdlandmask,
+    grdmask,
     grdpaste,
     grdproject,
     grdsample,

pygmt/src/__init__.py

@@ -25,6 +25,7 @@
 from pygmt.src.grdimage import grdimage
 from pygmt.src.grdinfo import grdinfo
 from pygmt.src.grdlandmask import grdlandmask
+from pygmt.src.grdmask import grdmask
 from pygmt.src.grdpaste import grdpaste
 from pygmt.src.grdproject import grdproject
 from pygmt.src.grdsample import grdsample

pygmt/src/grdmask.py

@@ -0,0 +1,172 @@
+"""
+grdmask - Create mask grid from polygons or point coverage.
+"""
+
+from collections.abc import Sequence
+from typing import Literal
+
+import xarray as xr
+from pygmt._typing import PathLike
+from pygmt.alias import Alias, AliasSystem
+from pygmt.clib import Session
+from pygmt.exceptions import GMTParameterError
+from pygmt.helpers import build_arg_list, fmt_docstring
+
+__doctest_skip__ = ["grdmask"]
+
+
+@fmt_docstring
+def grdmask(
+    data,
+    outgrid: PathLike | None = None,
+    spacing: Sequence[float | str] | None = None,
+    region: Sequence[float | str] | str | None = None,
+    outside: float | Literal["z", "id"] = 0,
+    edge: float | Literal["z", "id"] = 0,
+    inside: float | Literal["z", "id"] = 1,
+    verbose: Literal["quiet", "error", "warning", "timing", "info", "compat", "debug"]
+    | bool = False,
+    **kwargs,
+) -> xr.DataArray | None:
+    """
+    Create mask grid from polygons or point coverage.
+
+    Reads one or more files (or standard input) containing polygon or data point
+    coordinates, and creates a binary grid file where nodes that fall inside, on the
+    edge, or outside the polygons (or within the search radius from data points) are
+    assigned values based on ``outside``, ``edge``, and ``inside`` parameters.
+
+    The mask grid can be used to mask out specific regions in other grids using
+    :func:`pygmt.grdmath` or similar tools. For masking based on coastline features,
+    consider using :func:`pygmt.grdlandmask` instead.
+
+    Full GMT docs at :gmt-docs:`grdmask.html`.
+
+    **Aliases**
+
+    .. hlist::
+       :columns: 3
+
+       - G = outgrid
+       - I = spacing
+       - N = outside/edge/inside
+       - R = region
+       - V = verbose
+
+    Parameters
+    ----------
+    data
+        Pass in either a file name, :class:`pandas.DataFrame`, :class:`numpy.ndarray`,
+        or a list of file names containing the polygon(s) or data points. Input can be:
+
+        - **Polygon mode**: One or more files containing closed polygon coordinates
+        - **Point coverage mode**: Data points (used with ``search_radius`` parameter)
+    $outgrid
+    $spacing
+    outside
+        Set the value assigned to nodes outside the polygons. Default is 0.
+        Can be any number, or one of ``None``, ``"NaN"``, and ``np.nan`` for NaN.
+
+        When using ``inside="z"`` or ``inside="id"``, this sets the outside value
+        appended after the mode (e.g., ``outside=1, inside="z"`` gives ``-Nz/1``).
+    edge
+        Set the value assigned to nodes on the polygon edges. Default is 0.
+        Can be any number, or one of ``None``, ``"NaN"``, and ``np.nan`` for NaN.
+
+        When using ``inside="z"``, setting ``edge="z"`` treats edges as inside
+        (corresponds to ``-NZ``). Similarly, ``inside="id", edge="id"`` gives ``-NP``.
+        The combination ``inside="z", edge="id"`` or ``inside="id", edge="z"`` is
+        invalid and will raise an error.
+    inside
+        Set the value assigned to nodes inside the polygons. Default is 1.
+        Can be any number, or one of ``None``, ``"NaN"``, and ``np.nan`` for NaN.
+
+        Special values:
+
+        - ``"z"``: Use the z-value from polygon data (segment header ``-Zzval``,
+          ``-Lheader``, or via ``-aZ=name``). Corresponds to GMT ``-Nz``.
+        - ``"id"``: Use a running polygon ID number. Corresponds to GMT ``-Np``.
+    $region
+    $verbose
+
+    Returns
+    -------
+    ret
+        Return type depends on whether the ``outgrid`` parameter is set:
+
+        - :class:`xarray.DataArray` if ``outgrid`` is not set
+        - ``None`` if ``outgrid`` is set (grid output will be stored in the file set by
+          ``outgrid``)
+
+    Example
+    -------
+    >>> import pygmt
+    >>> import numpy as np
+    >>> # Create a simple polygon as a triangle
+    >>> polygon = np.array([[125, 30], [130, 30], [130, 35], [125, 30]])
+    >>> # Create a mask grid with 1 arc-degree spacing
+    >>> mask = pygmt.grdmask(data=polygon, spacing=1, region=[125, 130, 30, 35])
+    >>> mask.values
+    array([[0., 0., 0., 0., 0., 0.],
+           [0., 0., 0., 0., 1., 0.],
+           [0., 0., 1., 1., 1., 0.],
+           [0., 0., 1., 1., 1., 0.],
+           [0., 0., 1., 1., 1., 0.],
+           [0., 0., 0., 0., 0., 0.]])
+    """
+    if spacing is None or region is None:
+        raise GMTParameterError(required=["region", "spacing"])
+
+    # Build the -N parameter string
+    special_modes = {"z", "id"}
+    inside_is_special = inside in special_modes
+    edge_is_special = edge in special_modes
+
+    # Validate combinations
+    if inside_is_special and edge_is_special and inside != edge:
+        msg = f"Invalid combination: inside={inside!r} and edge={edge!r}. "
+        raise GMTParameterError(
+            reason=msg + "When both are special modes, they must be the same."
+        )
+
+    # Build -N argument
+    if inside_is_special:
+        # Mode: -Nz, -NZ, -Np, or -NP
+        mode_char = (
+            "Z"
+            if edge == inside
+            else "z"
+            if inside == "z"
+            else "P"
+            if edge == inside
+            else "p"
+        )
+        n_value = f"{mode_char}/{outside}" if outside != 0 else mode_char
+        aliasdict = AliasSystem(
+            I=Alias(spacing, name="spacing", sep="/", size=2),
+            N=n_value,
+        )
+    else:
+        # Standard mode: outside/edge/inside
+        aliasdict = AliasSystem(
+            I=Alias(spacing, name="spacing", sep="/", size=2),
+            N=Alias([outside, edge, inside], name="mask_values", sep="/", size=3),
+        )
+
+    aliasdict = aliasdict.add_common(
+        R=region,
+        V=verbose,
+    )
+    aliasdict.merge(kwargs)
+
+    with Session() as lib:
+        with (
+            lib.virtualfile_in(check_kind="vector", data=data) as vintbl,
+            lib.virtualfile_out(kind="grid", fname=outgrid) as voutgrd,
+        ):
+            aliasdict["G"] = voutgrd
+            lib.call_module(
+                module="grdmask",
+                args=build_arg_list(aliasdict, infile=vintbl),
+            )
+            return lib.virtualfile_to_raster(vfname=voutgrd, outgrid=outgrid)

pygmt/tests/test_grdmask.py

@@ -0,0 +1,123 @@
+"""
+Test pygmt.grdmask.
+"""
+
+from pathlib import Path
+
+import numpy as np
+import pytest
+import xarray as xr
+from pygmt import grdmask
+from pygmt.enums import GridRegistration, GridType
+from pygmt.exceptions import GMTParameterError
+from pygmt.helpers import GMTTempFile
+
+
+@pytest.fixture(scope="module", name="polygon_data")
+def fixture_polygon_data():
+    """
+    Create a simple polygon for testing.
+    A triangle polygon covering the region [125, 130, 30, 35].
+    """
+    return np.array([[125, 30], [130, 30], [130, 35], [125, 30]])
+
+
+@pytest.fixture(scope="module", name="expected_grid")
+def fixture_expected_grid():
+    """
+    Load the expected grdmask grid result.
+    """
+    return xr.DataArray(
+        data=[
+            [0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
+            [0.0, 0.0, 1.0, 1.0, 1.0, 0.0],
+            [0.0, 0.0, 0.0, 1.0, 1.0, 0.0],
+            [0.0, 0.0, 0.0, 0.0, 1.0, 0.0],
+            [0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
+            [0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
+        ],
+        coords={
+            "x": [125.0, 126.0, 127.0, 128.0, 129.0, 130.0],
+            "y": [30.0, 31.0, 32.0, 33.0, 34.0, 35.0],
+        },
+        dims=["y", "x"],
+    )
+
+
+def test_grdmask_outgrid(polygon_data, expected_grid):
+    """
+    Creates a mask grid with an outgrid argument.
+    """
+    with GMTTempFile(suffix=".nc") as tmpfile:
+        result = grdmask(
+            data=polygon_data,
+            outgrid=tmpfile.name,
+            spacing=1,
+            region=[125, 130, 30, 35],
+        )
+        assert result is None  # return value is None
+        assert Path(tmpfile.name).stat().st_size > 0  # check that outgrid exists
+        temp_grid = xr.load_dataarray(tmpfile.name, engine="gmt", raster_kind="grid")
+        # Check grid properties
+        assert temp_grid.dims == ("y", "x")
+        assert temp_grid.gmt.gtype is GridType.CARTESIAN
+        assert temp_grid.gmt.registration is GridRegistration.GRIDLINE
+        # Check grid values
+        xr.testing.assert_allclose(a=temp_grid, b=expected_grid)
+
+
+@pytest.mark.benchmark
+def test_grdmask_no_outgrid(polygon_data, expected_grid):
+    """
+    Test grdmask with no set outgrid.
+    """
+    result = grdmask(data=polygon_data, spacing=1, region=[125, 130, 30, 35])
+    # Check grid properties
+    assert isinstance(result, xr.DataArray)
+    assert result.dims == ("y", "x")
+    assert result.gmt.gtype is GridType.CARTESIAN
+    assert result.gmt.registration is GridRegistration.GRIDLINE
+    # Check grid values
+    xr.testing.assert_allclose(a=result, b=expected_grid)
+
+
+def test_grdmask_custom_mask_values(polygon_data):
+    """
+    Test grdmask with custom outside, edge, inside values.
+    """
+    result = grdmask(
+        data=polygon_data,
+        spacing=1,
+        region=[125, 130, 30, 35],
+        outside=10,
+        edge=20,
+        inside=30,
+    )
+    assert isinstance(result, xr.DataArray)
+    # Check that the grid has the right dimensions
+    assert result.shape == (6, 6)
+    # Check that we have values in the expected range
+    assert result.values.max() <= 30.0
+    assert result.values.min() >= 0.0
+
+
+def test_grdmask_fails():
+    """
+    Check that grdmask fails correctly when region and spacing are not given.
+    """
+    with pytest.raises(GMTParameterError):
+        grdmask(data=np.array([[0, 0], [1, 1], [1, 0], [0, 0]]))
+
+
+def test_grdmask_invalid_combination(polygon_data):
+    """
+    Check that grdmask fails when inside and edge have different special modes.
+    """
+    with pytest.raises(GMTParameterError):
+        grdmask(
+            data=polygon_data,
+            spacing=1,
+            region=[125, 130, 30, 35],
+            inside="z",
+            edge="id",
+        )