diff --git a/doc/api/index.rst b/doc/api/index.rst index 6fad08b8165..7c61b0f4738 100644 --- a/doc/api/index.rst +++ b/doc/api/index.rst @@ -126,6 +126,7 @@ Operations on tabular data blockmedian blockmode filter1d + grdmask nearneighbor project select diff --git a/pygmt/__init__.py b/pygmt/__init__.py index f10c4fad8a5..8adecdf429a 100644 --- a/pygmt/__init__.py +++ b/pygmt/__init__.py @@ -45,6 +45,7 @@ grdhisteq, grdinfo, grdlandmask, + grdmask, grdpaste, grdproject, grdsample, diff --git a/pygmt/src/__init__.py b/pygmt/src/__init__.py index 3c3fb2c39e3..b3a41b3f637 100644 --- a/pygmt/src/__init__.py +++ b/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 diff --git a/pygmt/src/grdmask.py b/pygmt/src/grdmask.py new file mode 100644 index 00000000000..9ae8dcf217a --- /dev/null +++ b/pygmt/src/grdmask.py @@ -0,0 +1,178 @@ +""" +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"] + + +def _alias_option_N( # noqa: N802 + outside: float | Literal["z", "id"] | None = None, + edge: float | Literal["z", "id"] | None = None, + inside: float | Literal["z", "id"] | None = None, +) -> Alias: + """ + Return an Alias object for the -N option. + + Builds the -N parameter string for grdmask based on the inside, edge, and + outside values. Handles special modes "z" (use z-value from polygon data) + and "id" (use running polygon ID). + + Examples + -------- + >>> _alias_option_N(outside=0, edge=0, inside=1)._value + '0/0/1' + >>> _alias_option_N(outside=1, edge=2, inside=3)._value + '1/2/3' + >>> _alias_option_N(outside=0, edge=0, inside="z")._value + 'z' + >>> _alias_option_N(outside=1, edge=0, inside="z")._value + 'z/1' + >>> _alias_option_N(outside=0, edge="z", inside="z")._value + 'Z' + >>> _alias_option_N(outside=0, edge=0, inside="id")._value + 'p' + >>> _alias_option_N(outside=0, edge="id", inside="id")._value + 'P' + """ + _inside_modes = {"z": "z", "id": "p"} + # Validate combinations + if inside in _inside_modes and edge in _inside_modes 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 in _inside_modes: + # Mode: -Nz, -NZ, -Np, or -NP + mode_char = _inside_modes[inside] + if edge == inside: + mode_char = mode_char.upper() + mask_values = mode_char if outside is None else [mode_char, outside] + else: + mask_values = [outside, edge, inside] + return Alias(mask_values, name="mask_values", sep="/", size=(2, 3)) + + +@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 + edge + inside + Set the value assigned to nodes outside, on the edge, or inside the polygons. + Can be any number, or one of ``None``, ``"NaN"``, and ``np.nan`` for NaN. + + ``inside`` can also be set to one of the following values: + + - ``"z"``: Use the z-value from polygon data (segment header ``-Zzval``, + ``-Lheader``, or via ``-aZ=name``). + - ``"id"``: Use a running polygon ID number. + + To treats edges as inside, using the same value as ``inside``. + $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"]) + + aliasdict = AliasSystem( + I=Alias(spacing, name="spacing", sep="/", size=2), + N=_alias_option_N(outside=outside, edge=edge, inside=inside), + ).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) diff --git a/pygmt/tests/test_grdmask.py b/pygmt/tests/test_grdmask.py new file mode 100644 index 00000000000..707869ba5e2 --- /dev/null +++ b/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", + )