pygmt.grdclip: Parameters 'between' and 'replace' now support 2-D sequences

seisman · seisman · commit 502bf6379a02 · 2025-03-31T14:12:41.000+08:00
diff --git a/pygmt/src/grdclip.py b/pygmt/src/grdclip.py
@@ -2,48 +2,129 @@
 grdclip - Clip the range of grid values.
 """
 
+from collections.abc import Sequence
+
 import xarray as xr
 from pygmt.clib import Session
+from pygmt.exceptions import GMTInvalidInput
 from pygmt.helpers import (
     build_arg_list,
     deprecate_parameter,
     fmt_docstring,
+    is_nonstr_iter,
     kwargs_to_strings,
     use_alias,
 )
 
 __doctest_skip__ = ["grdclip"]
 
 
+def _parse_sequence(name, value, separator="/", size=2, ndim=1):
+    """
+    Parse a 1-D or 2-D sequence of values and join them by a separator.
+
+    Parameters
+    ----------
+    name
+        The parameter name.
+    value
+        The 1-D or 2-D sequence of values to parse.
+    separator
+        The separator to join the values.
+    size
+        The number of values in the sequence.
+    ndim
+        The expected maximum number of dimensions of the sequence.
+
+    Returns
+    -------
+    str
+        The parsed sequence.
+
+    Examples
+    --------
+    >>> _parse_sequence("above_or_below", [1000, 0], size=2, ndim=1)
+    '1000/0'
+    >>> _parse_sequence("between", [1000, 1500, 10000], size=3, ndim=2)
+    '1000/1500/10000'
+    >>> _parse_sequence("between", [[1000, 1500, 10000]], size=3, ndim=2)
+    ['1000/1500/10000']
+    >>> _parse_sequence(
+    ...     "between", [[1000, 1500, 10000], [1500, 2000, 20000]], size=3, ndim=2
+    ... )
+    ['1000/1500/10000', '1500/2000/20000']
+    >>> _parse_sequence("replace", [1000, 0], size=2, ndim=2)
+    '1000/0'
+    >>> _parse_sequence("replace", [[1000, 0]], size=2, ndim=2)
+    ['1000/0']
+    >>> _parse_sequence("replace", [[1000, 0], [1500, 10000]], size=2, ndim=2)
+    ['1000/0', '1500/10000']
+    >>> _parse_sequence("any", "1000/100")
+    '1000/100'
+    >>> _parse_sequence("any", None)
+    >>> _parse_sequence("any", [])
+    []
+    >>> _parse_sequence("above_or_below", [[100, 1000], [1500, 2000]], size=2, ndim=1)
+    Traceback (most recent call last):
+        ...
+    pygmt.exceptions.GMTInvalidInput: Parameter ... must be a 1-D sequence...
+    """
+    # Return the value as is if not a sequence (e.g., str or None) or empty.
+    if not is_nonstr_iter(value) or len(value) == 0:
+        return value
+
+    # 1-D sequence
+    if not is_nonstr_iter(value[0]):
+        if len(value) != size:
+            msg = (
+                f"Parameter '{name}' must be a 1-D sequence of {size} values, "
+                f"but got {len(value)} values."
+            )
+            raise GMTInvalidInput(msg)
+        return separator.join(str(i) for i in value)
+
+    # 2-D sequence
+    if ndim == 1:
+        msg = f"Parameter '{name}' must be a 1-D sequence, not a 2-D sequence."
+        raise GMTInvalidInput(msg)
+
+    if any(len(i) != size for i in value):
+        msg = (
+            f"Parameter '{name}' must be a 2-D sequence with each sub-sequence "
+            f"having {size} values."
+        )
+        raise GMTInvalidInput(msg)
+    return [separator.join(str(j) for j in value[i]) for i in range(len(value))]
+
+
 # TODO(PyGMT>=0.19.0): Remove the deprecated "new" parameter.
 @fmt_docstring
 @deprecate_parameter("new", "replace", "v0.15.0", remove_version="v0.19.0")
-@use_alias(
-    R="region",
-    Sa="above",
-    Sb="below",
-    Si="between",
-    Sr="replace",
-    V="verbose",
-)
-@kwargs_to_strings(
-    R="sequence",
-    Sa="sequence",
-    Sb="sequence",
-    Si="sequence",
-    Sr="sequence",
-)
-def grdclip(grid, outgrid: str | None = None, **kwargs) -> xr.DataArray | None:
-    r"""
+@use_alias(R="region", V="verbose")
+@kwargs_to_strings(R="sequence")
+def grdclip(
+    grid,
+    outgrid: str | None = None,
+    above: Sequence[float] | None = None,
+    below: Sequence[float] | None = None,
+    between: Sequence[float] | Sequence[Sequence[float]] | None = None,
+    replace: Sequence[float] | Sequence[Sequence[float]] | None = None,
+    **kwargs,
+) -> xr.DataArray | None:
+    """
     Clip the range of grid values.
 
-    Produce a clipped ``outgrid`` or :class:`xarray.DataArray` version of the
-    input ``grid`` file.
+    This function operates on the values of a grid. It can:
+
+    - Set values smaller than a threshold to a new value
+    - Set values larger than a threshold to a new value
+    - Set values within a range to a new value
+    - Replace individual values with a new value
 
-    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.
+    Such operations are useful when you want all of a continent or an ocean to fall into
+    one color or gray shade in image processing, when clipping of the range of data
+    values is required, or for reclassification of data values. The values can be any
+    number or even NaN (Not a Number).
 
     Full option list at :gmt-docs:`grdclip.html`
 
@@ -54,19 +135,23 @@ def grdclip(grid, outgrid: str | None = None, **kwargs) -> xr.DataArray | None:
     {grid}
     {outgrid}
     {region}
-    above : str or list
-        [*high*, *above*].
-        Set all data[i] > *high* to *above*.
-    below : str or list
-        [*low*, *below*].
-        Set all data[i] < *low* to *below*.
-    between : str or list
-        [*low*, *high*, *between*].
-        Set all data[i] >= *low* and <= *high* to *between*.
-    replace : str or list
-        [*old*, *new*].
-        Set all data[i] == *old* to *new*. This is mostly useful when
-        your data are known to be integer values.
+    above
+        Pass a sequence of two values in the form of (*high*, *above*), to set all node
+        values greater than *high* to *above*.
+    below
+        Pass a sequence of two values in the form of (*low*, *below*) to set all node
+        values less than *low* to *below*.
+    between
+        Pass a sequence of three values in the form of (*low*, *high*, *between*) to set
+        all node values between *low* and *high* to *between*. It can also accept a
+        sequence of sequences (e.g., list of lists or 2-D numpy array) to set different
+        values for different ranges.
+    replace
+        Pass a sequence of two values in the form of (*old*, *new*) to replace all node
+        values equal to *old* with *new*. It can also accept a sequence of sequences
+        (e.g., list of lists or 2-D numpy array) to replace different old values with
+        different new values. This is mostly useful when your data are known to be
+        integer values.
     {verbose}
 
     Returns
@@ -96,6 +181,19 @@ def grdclip(grid, outgrid: str | None = None, **kwargs) -> xr.DataArray | None:
     >>> [new_grid.data.min(), new_grid.data.max()]
     [0.0, 10000.0]
     """
+    if all(v is None for v in (above, below, between, replace)):
+        msg = (
+            "Must specify at least one of the following parameters: ",
+            "'above', 'below', 'between', or 'replace'.",
+        )
+        raise GMTInvalidInput(msg)
+
+    # Parse the -S option.
+    kwargs["Sa"] = _parse_sequence("above", above, size=2, ndim=1)
+    kwargs["Sb"] = _parse_sequence("below", below, size=2, ndim=1)
+    kwargs["Si"] = _parse_sequence("between", between, size=3, ndim=2)
+    kwargs["Sr"] = _parse_sequence("replace", replace, size=2, ndim=2)
+
     with Session() as lib:
         with (
             lib.virtualfile_in(check_kind="raster", data=grid) as vingrd,
diff --git a/pygmt/tests/test_grdclip.py b/pygmt/tests/test_grdclip.py
@@ -88,3 +88,18 @@ def test_grdclip_replace():
     with pytest.warns(FutureWarning):
         grid = grdclip(grid=grid, new=[1, 3])  # Replace 1 with 3
         npt.assert_array_equal(np.unique(grid), [2, 3])
+
+
+def test_grdclip_between_repeated():
+    """
+    Test passing a 2D sequence to the between parameter for grdclip.
+    """
+    grid = load_static_earth_relief()
+    # Replace values in the range 0-250 with 0, 250-500 with 1, 500-750 with 2, and
+    # 750-1000 with 3
+    result = grdclip(
+        grid,
+        between=[[0, 250, 0], [250, 500, 1], [500, 750, 2], [750, 1000, 3]],
+    )
+    # result should have 4 unique values.
+    npt.assert_array_equal(np.unique(result.data), [0, 1, 2, 3])
