Wrap grdpaste for joining two grids along their common edge (#4399)

Co-authored-by: Dongdong Tian <seisman.info@gmail.com>

Author: Chuan1937

doc/api/index.rst

@@ -155,6 +155,7 @@ Operations on raster data
     grdhisteq.equalize_grid
     grdhisteq.compute_bins
     grdlandmask
+    grdpaste
     grdproject
     grdsample
     grdtrack

pygmt/__init__.py

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

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.grdpaste import grdpaste
 from pygmt.src.grdproject import grdproject
 from pygmt.src.grdsample import grdsample
 from pygmt.src.grdtrack import grdtrack

pygmt/src/grdpaste.py

@@ -0,0 +1,90 @@
+"""
+grdpaste - Join two grids along their common edge.
+"""
+
+from typing import Literal
+
+import xarray as xr
+from pygmt._typing import PathLike
+from pygmt.alias import AliasSystem
+from pygmt.clib import Session
+from pygmt.exceptions import GMTTypeError
+from pygmt.helpers import build_arg_list, data_kind, fmt_docstring, use_alias
+
+
+@fmt_docstring
+@use_alias(f="coltypes")
+def grdpaste(
+    grid1: PathLike | xr.DataArray,
+    grid2: PathLike | xr.DataArray,
+    outgrid: PathLike | None = None,
+    verbose: Literal["quiet", "error", "warning", "timing", "info", "compat", "debug"]
+    | bool = False,
+    **kwargs,
+) -> xr.DataArray | None:
+    """
+    Join two grids along their common edge.
+
+    Combine ``grid1`` and ``grid2`` into a single grid by pasting them together along
+    their common edge. The two input grids must have the same grid spacings and
+    registration, and must have one edge in common. If in doubt, check with
+    :func:`pygmt.grdinfo` and use :func:`pygmt.grdcut` and/or :func:`pygmt.grdsample` if
+    necessary to prepare the edge joint. Note: For geographical grids, you may have to
+    use ``coltypes`` to handle periodic longitudes unless the input grids are properly
+    recognized as such via their meta-data. For stitching multiple grids, see
+    ``grdblend`` (not implemented in PyGMT yet) instead.
+
+    Full GMT docs at :gmt-docs:`grdpaste.html`.
+
+    $aliases
+       - G = outgrid
+       - V = verbose
+
+    Parameters
+    ----------
+    grid1
+    grid2
+        The two grids to be pasted. Accepts either file names or
+        :class:`xarray.DataArray` objects.
+
+        **Notes:**
+
+        #. Both grids must be of the same type—mixing a file name with an
+           ``xarray.DataArray`` is not allowed.
+        #. Passing two ``xarray.DataArray`` objects requires GMT>6.6.0 due to
+           an upstream GMT bug.
+    $outgrid
+    $verbose
+    $coltypes
+
+    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``)
+    """
+    # Check if grid1 and grid2 are of the same kind
+    if data_kind(grid1) != data_kind(grid2):
+        raise GMTTypeError(
+            (type(grid1), type(grid2)),
+            reason="Both input grids must be of the same type (file or xarray.DataArray).",
+        )
+
+    aliasdict = AliasSystem().add_common(V=verbose)
+    aliasdict.merge(kwargs)
+
+    with Session() as lib:
+        with (
+            lib.virtualfile_in(check_kind="raster", data=grid1) as vingrd1,
+            lib.virtualfile_in(check_kind="raster", data=grid2) as vingrd2,
+            lib.virtualfile_out(kind="grid", fname=outgrid) as voutgrd,
+        ):
+            aliasdict["G"] = voutgrd
+            lib.call_module(
+                module="grdpaste",
+                args=build_arg_list(aliasdict, infile=[vingrd1, vingrd2]),
+            )
+            return lib.virtualfile_to_raster(vfname=voutgrd, outgrid=outgrid)

pygmt/tests/test_grdpaste.py

@@ -0,0 +1,128 @@
+"""
+Test pygmt.grdpaste.
+"""
+
+import pytest
+import xarray as xr
+from packaging.version import Version
+from pygmt import grdcut, grdpaste
+from pygmt.clib import __gmt_version__
+from pygmt.exceptions import GMTTypeError
+from pygmt.helpers import GMTTempFile
+from pygmt.helpers.testing import load_static_earth_relief
+
+
+@pytest.fixture(scope="module", name="grid")
+def fixture_grid():
+    """
+    Load the grid data from the sample earth_relief file.
+    """
+    return load_static_earth_relief()
+
+
+@pytest.fixture(scope="module", name="grid_top")
+def fixture_grid_top(grid):
+    """
+    Load the top part of the grid data from the sample earth_relief file.
+    """
+    return grdcut(grid, region=[-53, -49, -19, -16])
+
+
+@pytest.fixture(scope="module", name="grid_bottom")
+def fixture_grid_bottom(grid):
+    """
+    Load the bottom part of the grid data from the sample earth_relief file.
+    """
+    return grdcut(grid, region=[-53, -49, -22, -19])
+
+
+def test_grdpaste_file_in_file_out(grid):
+    """
+    Test grdpaste with file input and file output.
+    """
+    with (
+        GMTTempFile(suffix=".nc") as tmp1,
+        GMTTempFile(suffix=".nc") as tmp2,
+        GMTTempFile(suffix=".nc") as tmpout,
+    ):
+        grdcut(grid, region=[-53, -49, -19, -16], outgrid=tmp1.name)
+        grdcut(grid, region=[-53, -49, -22, -19], outgrid=tmp2.name)
+        result = grdpaste(grid1=tmp1.name, grid2=tmp2.name, outgrid=tmpout.name)
+        assert result is None  # grdpaste returns None if output to a file
+        temp_grid = xr.load_dataarray(tmpout.name, engine="gmt", raster_kind="grid")
+        assert isinstance(temp_grid, xr.DataArray)
+        assert temp_grid.shape == (6, 4)
+        # Check that the result has the expected min and max values
+        assert temp_grid.min().values == 345.5
+        assert temp_grid.max().values == 886.0
+
+
+def test_grdpaste_file_in_xarray_out(grid):
+    """
+    Test grdpaste with file input and xarray output.
+    """
+    with (
+        GMTTempFile(suffix=".nc") as tmp1,
+        GMTTempFile(suffix=".nc") as tmp2,
+    ):
+        grdcut(grid, region=[-53, -49, -19, -16], outgrid=tmp1.name)
+        grdcut(grid, region=[-53, -49, -22, -19], outgrid=tmp2.name)
+        result = grdpaste(grid1=tmp1.name, grid2=tmp2.name)
+        assert isinstance(result, xr.DataArray)
+        assert result.shape == (6, 4)
+        # Check that the result has the expected min and max values
+        assert result.min().values == 345.5
+        assert result.max().values == 886.0
+
+
+# TODO(GMT>6.6.0): Remove the xfail marker.
+@pytest.mark.xfail(
+    condition=Version(__gmt_version__) <= Version("6.6.0"),
+    reason="Upstream bug fixed in https://github.com/GenericMappingTools/gmt/pull/8901",
+)
+def test_grdpaste(grid_top, grid_bottom):
+    """
+    Test grdpaste by pasting two grids together along their common edge.
+    """
+    # Paste the two grids back together
+    result = grdpaste(grid1=grid_top, grid2=grid_bottom)
+    # Check that the result is a DataArray
+    assert isinstance(result, xr.DataArray)
+    # Check that the result has the expected shape
+    # grid_top has 3x4, grid_bottom has 3x4, so result should have 6x4
+    assert result.shape == (6, 4)
+    # Check that the result has the expected min and max values
+    assert result.min().values == 345.5
+    assert result.max().values == 886.0
+
+
+# TODO(GMT>6.6.0): Remove the xfail marker.
+@pytest.mark.xfail(
+    condition=Version(__gmt_version__) <= Version("6.6.0"),
+    reason="Upstream bug fixed in https://github.com/GenericMappingTools/gmt/pull/8901",
+)
+def test_grdpaste_outgrid(grid_top, grid_bottom):
+    """
+    Test grdpaste with outgrid parameter.
+    """
+    # Paste the two grids back together and save to file
+    with GMTTempFile(suffix=".nc") as tmpfile:
+        result = grdpaste(grid1=grid_top, grid2=grid_bottom, outgrid=tmpfile.name)
+        assert result is None  # grdpaste returns None if output to a file
+        temp_grid = xr.load_dataarray(tmpfile.name, engine="gmt", raster_kind="grid")
+        assert isinstance(temp_grid, xr.DataArray)
+        assert temp_grid.shape == (6, 4)
+        # Check that the result has the expected min and max values
+        assert temp_grid.min().values == 345.5
+        assert temp_grid.max().values == 886.0
+
+
+def test_grdpaste_mixed_inputs_file_xarray(grid, grid_bottom):
+    """
+    Test that mixing file and xarray inputs raises GMTTypeError.
+    """
+    with GMTTempFile(suffix=".nc") as tmp1:
+        grdcut(grid, region=[-53, -49, -19, -16], outgrid=tmp1.name)
+        # This should raise GMTTypeError because we're mixing file and xarray inputs
+        with pytest.raises(GMTTypeError):
+            grdpaste(grid1=tmp1.name, grid2=grid_bottom)