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,110 @@
+"""
+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,
+    mask_values: Sequence[float] | None = None,
+    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 ``mask_values``.
+
+    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
+       - G = outgrid
+       - I = spacing
+       - N = mask_values
+       - 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
+    mask_values
+        Set the values that will be assigned to nodes. Provide three values in the form
+        [*outside*, *edge*, *inside*]. Default is ``[0, 0, 1]``, meaning nodes outside
+        and on the edge are set to 0, and nodes inside are set to 1.
+
+        Values can be any number, or one of ``None``, ``"NaN"``, and ``np.nan`` for
+        setting nodes to NaN.
+    $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])
+    """
+    if spacing is None or region is None:
+        raise GMTParameterError(required=["region", "spacing"])
+
+    aliasdict = AliasSystem(
+        I=Alias(spacing, name="spacing", sep="/", size=2),
+        N=Alias(mask_values, name="mask_values", sep="/", size=3),
+    ).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,107 @@
+"""
+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 mask_values.
+    """
+    result = grdmask(
+        data=polygon_data,
+        spacing=1,
+        region=[125, 130, 30, 35],
+        mask_values=[10, 20, 30],  # outside, edge, inside
+    )
+    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]]))