From c7330da501ab1817f247451edd45937bc176754a Mon Sep 17 00:00:00 2001
From: Will Schlitzer <schlitzer90@gmail.com>
Date: Tue, 18 May 2021 07:55:15 +0100
Subject: [PATCH] Wrap grdclip function (#1261)

Wrap the gmt function grdclip and add tests for it.

Co-authored-by: Wei Ji <23487320+weiji14@users.noreply.github.com>
Co-authored-by: Dongdong Tian <seisman.info@gmail.com>
Co-authored-by: Wei Ji <23487320+weiji14@users.noreply.github.com>
---
 doc/api/index.rst           |  1 +
 pygmt/__init__.py           |  1 +
 pygmt/src/__init__.py       |  1 +
 pygmt/src/grdclip.py        | 89 +++++++++++++++++++++++++++++++++++++
 pygmt/tests/test_grdclip.py | 47 ++++++++++++++++++++
 5 files changed, 139 insertions(+)
 create mode 100644 pygmt/src/grdclip.py
 create mode 100644 pygmt/tests/test_grdclip.py

diff --git a/doc/api/index.rst b/doc/api/index.rst
index b6cb79a26d6..dec563a4b8d 100644
--- a/doc/api/index.rst
+++ b/doc/api/index.rst
@@ -88,6 +88,7 @@ Operations on grids:
 .. autosummary::
     :toctree: generated
 
+    grdclip
     grdcut
     grdfilter
     grdtrack
diff --git a/pygmt/__init__.py b/pygmt/__init__.py
index 1dc2ed79ac4..f9a7d00e1b0 100644
--- a/pygmt/__init__.py
+++ b/pygmt/__init__.py
@@ -33,6 +33,7 @@
     blockmedian,
     config,
     grd2cpt,
+    grdclip,
     grdcut,
     grdfilter,
     grdinfo,
diff --git a/pygmt/src/__init__.py b/pygmt/src/__init__.py
index 083a2970109..5189c096243 100644
--- a/pygmt/src/__init__.py
+++ b/pygmt/src/__init__.py
@@ -10,6 +10,7 @@
 from pygmt.src.config import config
 from pygmt.src.contour import contour
 from pygmt.src.grd2cpt import grd2cpt
+from pygmt.src.grdclip import grdclip
 from pygmt.src.grdcontour import grdcontour
 from pygmt.src.grdcut import grdcut
 from pygmt.src.grdfilter import grdfilter
diff --git a/pygmt/src/grdclip.py b/pygmt/src/grdclip.py
new file mode 100644
index 00000000000..f0b5dbd7004
--- /dev/null
+++ b/pygmt/src/grdclip.py
@@ -0,0 +1,89 @@
+"""
+grdclip - Change the range and extremes of grid values.
+"""
+
+import xarray as xr
+from pygmt.clib import Session
+from pygmt.helpers import (
+    GMTTempFile,
+    build_arg_string,
+    fmt_docstring,
+    kwargs_to_strings,
+    use_alias,
+)
+
+
+@fmt_docstring
+@use_alias(
+    G="outgrid",
+    R="region",
+    Sa="above",
+    Sb="below",
+    Si="between",
+    Sr="new",
+    V="verbose",
+)
+@kwargs_to_strings(
+    R="sequence",
+    Sa="sequence",
+    Sb="sequence",
+    Si="sequence",
+    Sr="sequence",
+)
+def grdclip(grid, **kwargs):
+    r"""
+    Sets values in a grid that meet certain criteria to a new value.
+
+    Produce a clipped ``outgrid`` or :class:`xarray.DataArray` version of the
+    input ``grid`` file.
+
+    The parameters ``above`` and ``below`` allow for a given value to be set
+    for values above or below a set amount, respectively. This allows for
+    extreme values in a grid, such as points below a certain depth when
+    plotting Earth relief, to all be set to the same value.
+
+    Full option list at :gmt-docs:`grdclip.html`
+
+    {aliases}
+
+    Parameters
+    ----------
+    grid : str or xarray.DataArray
+        The file name of the input grid or the grid loaded as a DataArray.
+    outgrid : str or None
+        The name of the output netCDF file with extension .nc to store the grid
+        in.
+    {R}
+    above : str or list or tuple
+        [*high*, *above*].
+        Set all data[i] > *high* to *above*.
+    below : str or list or tuple
+        [*low*, *below*].
+        Set all data[i] < *low* to *below*.
+    between : str or list or tuple
+        [*low*, *high*, *between*].
+        Set all data[i] >= *low* and <= *high* to *between*.
+    new : str or list or tuple
+        [*old*, *new*].
+        Set all data[i] == *old* to *new*. This is mostly useful when
+        your data are known to be integer values.
+    {V}
+    """
+    with GMTTempFile(suffix=".nc") as tmpfile:
+        with Session() as lib:
+            file_context = lib.virtualfile_from_data(check_kind="raster", data=grid)
+            with file_context as infile:
+                if "G" not in kwargs.keys():  # if outgrid is unset, output to tempfile
+                    kwargs.update({"G": tmpfile.name})
+                outgrid = kwargs["G"]
+                arg_str = " ".join([infile, build_arg_string(kwargs)])
+                lib.call_module("grdclip", arg_str)
+
+        if outgrid == tmpfile.name:  # if user did not set outgrid, return DataArray
+            with xr.open_dataarray(outgrid) as dataarray:
+                result = dataarray.load()
+                _ = result.gmt  # load GMTDataArray accessor information
+        else:
+            result = None  # if user sets an outgrid, return None
+
+        return result
diff --git a/pygmt/tests/test_grdclip.py b/pygmt/tests/test_grdclip.py
new file mode 100644
index 00000000000..af08406eee2
--- /dev/null
+++ b/pygmt/tests/test_grdclip.py
@@ -0,0 +1,47 @@
+"""
+Tests for grdclip.
+"""
+import os
+
+import numpy.testing as npt
+import pytest
+from pygmt import grdclip, grdinfo
+from pygmt.datasets import load_earth_relief
+from pygmt.helpers import GMTTempFile
+
+
+@pytest.fixture(scope="module", name="grid")
+def fixture_grid():
+    """
+    Load the grid data from the sample earth_relief file.
+    """
+    return load_earth_relief(resolution="01d", region=[-5, 5, -5, 5])
+
+
+def test_grdclip_outgrid(grid):
+    """
+    Test the below and above parameters for grdclip and creates a test outgrid.
+    """
+    with GMTTempFile(suffix=".nc") as tmpfile:
+        result = grdclip(
+            grid=grid, outgrid=tmpfile.name, below=[-1500, -1800], above=[-200, 40]
+        )
+        assert result is None  # return value is None
+        assert os.path.exists(path=tmpfile.name)  # check that outgrid exists
+        result = (
+            grdinfo(grid=tmpfile.name, force_scan=0, per_column="n").strip().split()
+        )
+    assert int(result[4]) == -1800  # minimum value
+    assert int(result[5]) == 40  # maximum value
+
+
+def test_grdclip_no_outgrid(grid):
+    """
+    Test the below and above parameters for grdclip with no set outgrid.
+    """
+    temp_grid = grdclip(grid=grid, below=[-1500, -1800], above=[-200, 40])
+    assert temp_grid.dims == ("lat", "lon")
+    assert temp_grid.gmt.gtype == 1  # Geographic grid
+    assert temp_grid.gmt.registration == 1  # Pixel registration
+    npt.assert_allclose(temp_grid.min(), -1800)
+    npt.assert_allclose(temp_grid.max(), 40)