pygmt.grdclip: Parameters 'between' and 'new' now can accept a 2-D sequence

Author: seisman

pygmt/src/grdclip.py

@@ -2,40 +2,93 @@
 grdclip - Clip the range of grid values.
 """
 
+from collections.abc import Sequence
+
 import xarray as xr
 from pygmt.clib import Session
-from pygmt.helpers import build_arg_list, fmt_docstring, kwargs_to_strings, use_alias
+from pygmt.exceptions import GMTInvalidInput
+from pygmt.helpers import (
+    build_arg_list,
+    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 sequence of values from a string or a list.
+
+    >>> _parse_sequence("above", [1000, 0], size=2, ndim=1)
+    '1000/0'
+    >>> _parse_sequence("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=1)
+    '1000/0'
+    >>> _parse_sequence("replace", [[1000, 0], [1500, 10000]], size=2, ndim=2)
+    ['1000/0', '1500/10000']
+    """
+    if not is_nonstr_iter(value):  # Not a sequence. Likely str or None.
+        return value
+
+    # A sequence of sequences.
+    if len(value) == 0:
+        return None
+    if is_nonstr_iter(value[0]):  # 2-D sequence
+        if ndim == 1:
+            msg = f"Parameter '{name}' must be a 1-D sequence, not a 2-D sequence."
+            raise GMTInvalidInput(msg)
+
+        actual_sizes = {len(i) for i in value}
+        if len(actual_sizes) != 1 or actual_sizes != {size}:
+            msg = f"Parameter '{name}' must be a 1-D or 2D sequence with {size} values."
+            raise GMTInvalidInput(msg)
+        return [separator.join(str(j) for j in value[i]) for i in range(len(value))]
+
+    # A sequence.
+    if len(value) != size:
+        msg = f"Parameter '{name}' must be a 1-D sequence of {size} values, but got {len(value)}."
+        raise GMTInvalidInput(msg)
+    return separator.join(str(i) for i in value)
+
+
 @fmt_docstring
-@use_alias(
-    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, outgrid: str | None = None, **kwargs) -> xr.DataArray | None:
+@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,
+    new: Sequence[float] | Sequence[Sequence[float]] | None = None,
+    **kwargs,
+) -> xr.DataArray | None:
     r"""
     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`
 
@@ -46,19 +99,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*.
-    new : 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.
+    new
+        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
@@ -88,6 +145,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, new)):
+        msg = (
+            "Must specify at least one of the following parameters: ",
+            "'above', 'below', 'between', or 'new'.",
+        )
+        raise GMTInvalidInput(msg)
+
+    # Parse the -S option.
+    kwargs["Sa"] = _parse_sequence("above", above, size=2)
+    kwargs["Sb"] = _parse_sequence("below", below, size=2)
+    kwargs["Si"] = _parse_sequence("between", between, size=3, ndim=2)
+    kwargs["Sr"] = _parse_sequence("new", new, size=2, ndim=2)
+
     with Session() as lib:
         with (
             lib.virtualfile_in(check_kind="raster", data=grid) as vingrd,

pygmt/tests/test_grdclip.py

@@ -4,6 +4,8 @@
 
 from pathlib import Path
 
+import numpy as np
+import numpy.testing as npt
 import pytest
 import xarray as xr
 from pygmt import grdclip, load_dataarray
@@ -69,3 +71,18 @@ def test_grdclip_no_outgrid(grid, expected_grid):
     assert temp_grid.gmt.gtype == GridType.GEOGRAPHIC
     assert temp_grid.gmt.registration == GridRegistration.PIXEL
     xr.testing.assert_allclose(a=temp_grid, b=expected_grid)
+
+
+def test_grdclip_between():
+    """
+    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])