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..0a3e0867a48 --- /dev/null +++ b/pygmt/src/grdmask.py @@ -0,0 +1,193 @@ +""" +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 | 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()._value + >>> _alias_option_N(outside=1, edge=2, inside=3)._value + '1/2/3' + >>> _alias_option_N(outside=3)._value + '3/0/1' + >>> _alias_option_N(inside="z")._value + 'z' + >>> _alias_option_N(outside=1, inside="z")._value + 'z/1' + >>> _alias_option_N(edge="z", inside="z")._value + 'Z' + >>> _alias_option_N(inside="id")._value + 'p' + >>> _alias_option_N(edge="id", inside="id")._value + 'P' + """ + _inside_modes = {"z": "z", "id": "p"} + # Validate combinations + if edge in _inside_modes and inside != edge: + msg = f"Invalid combination: edge={edge!r} requires inside={edge!r}." + raise GMTParameterError(reason=msg) + + 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 = _inside_modes[inside] # type: ignore[index] + if edge == inside: + mode = mode.upper() + mask_values: str | list[str | float] | None = ( + mode if outside is None else [mode, outside] # type: ignore[assignment] + ) + elif inside is None and outside is None and edge is None: + # All three are None, return None (GMT uses default 0/0/1) + mask_values = None + else: + # Build the full mask with defaults for any missing values + mask_values = [ + 0 if outside is None else outside, + 0 if edge is None else edge, + 1 if inside is None else inside, + ] # type: ignore[list-item] + 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 | None = None, + edge: float | Literal["z", "id"] | None = None, + inside: float | Literal["z", "id"] | 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 containing polygon or data point + coordinates, and creates a grid 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 to an ASCII data table, a 2-D $table_classes + containg 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 treat edges as inside, use 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., 1., 1., 1., 0.], + [0., 0., 0., 1., 1., 0.], + [0., 0., 0., 0., 1., 0.], + [0., 0., 0., 0., 0., 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..748d0704b90 --- /dev/null +++ b/pygmt/tests/test_grdmask.py @@ -0,0 +1,181 @@ +""" +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"], + ) + + +@pytest.fixture(scope="module", name="expected_grid_outside_only") +def fixture_expected_grid_outside_only(): + """ + Load the expected grdmask grid result when only outside is set. + """ + return xr.DataArray( + data=[ + [0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + [3.0, 0.0, 1.0, 1.0, 1.0, 0.0], + [3.0, 3.0, 0.0, 1.0, 1.0, 0.0], + [3.0, 3.0, 3.0, 0.0, 1.0, 0.0], + [3.0, 3.0, 3.0, 3.0, 0.0, 0.0], + [3.0, 3.0, 3.0, 3.0, 3.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_outside_only(polygon_data, expected_grid_outside_only): + """ + Test grdmask when only outside is set. + """ + result = grdmask( + data=polygon_data, + spacing=1, + region=[125, 130, 30, 35], + outside=3, + ) + + assert isinstance(result, xr.DataArray) + assert result.dims == ("y", "x") + assert result.gmt.gtype is GridType.CARTESIAN + assert result.gmt.registration is GridRegistration.GRIDLINE + xr.testing.assert_allclose(a=result, b=expected_grid_outside_only) + + +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", + ) + + +@pytest.mark.parametrize( + ("edge", "inside"), + [("z", None), ("id", None), ("z", 5), ("id", 5)], +) +def test_grdmask_invalid_edge_special_mode(polygon_data, edge, inside): + """ + Check that special edge modes require the same special inside mode. + """ + with pytest.raises(GMTParameterError): + grdmask( + data=polygon_data, + spacing=1, + region=[125, 130, 30, 35], + edge=edge, + inside=inside, + )